add functional tests (#38)

* add functional tests

* fix incorrect loopback address

* fix assertion

* lint: ignore funlen for NewRootCommand

* address linter warnings

Author: bschaatsbergen

.golangci.yaml

@@ -8,6 +8,10 @@ issues:
       linters:
         - funlen
 
+    - source: "^func NewRootCommand"
+      linters:
+        - funlen
+
 linters:
   disable-all: true
   enable:

cmd/root.go

@@ -21,16 +21,16 @@ const (
 )
 
 var (
-	version string
-
+	version = "dev"
 	debug   bool
 	json    bool
 	noColor bool
+	server  string
+	qtype   string
+)
 
-	server string
-	qtype  string
-
-	rootCmd = &cobra.Command{
+func NewRootCommand() *cobra.Command {
+	cmd := &cobra.Command{
 		Use:   "zns",
 		Short: "zns is a command-line utility for querying DNS records and displaying them in human- or machine-readable formats.",
 		Long:  "zns is a command-line utility for querying DNS records, displaying them in a human-readable, colored format that includes type, name, TTL, and value. It supports various DNS record types, concurrent queries for improved performance, JSON output format for machine-readable results, and options to write output to a file or query a specific DNS server.",
@@ -138,7 +138,13 @@ var (
 				}
 			}
 
-			querier := query.NewQueryClient(fmt.Sprintf("%s:53", server), new(dns.Client), logger)
+			// If the server address does not already include a port,
+			// append the default DNS port (53) to it.
+			if !strings.Contains(server, ":") {
+				server = fmt.Sprintf("%s:53", server)
+			}
+
+			querier := query.NewQueryClient(server, new(dns.Client), logger)
 
 			logger.Debug("Creating querier", "server", server, "qtype", qtype, "domain", args[0])
 
@@ -181,17 +187,18 @@ var (
 			return nil
 		},
 	}
-)
 
-func init() {
-	rootCmd.CompletionOptions.DisableDefaultCmd = true
-	rootCmd.Flags().StringVarP(&server, "server", "s", "", "DNS server to query")
-	rootCmd.Flags().StringVarP(&qtype, "query-type", "q", "", "DNS query type")
-	rootCmd.Flags().BoolVar(&debug, "debug", false, "If set, debug output is printed")
-	rootCmd.Flags().BoolVar(&json, "json", false, "If set, output is printed in JSON format.")
+	cmd.CompletionOptions.DisableDefaultCmd = true
+	cmd.Flags().StringVarP(&server, "server", "s", "", "DNS server to query")
+	cmd.Flags().StringVarP(&qtype, "query-type", "q", "", "DNS query type")
+	cmd.Flags().BoolVar(&debug, "debug", false, "Enable debug output")
+	cmd.Flags().BoolVar(&json, "json", false, "Output in JSON format")
+
+	return cmd
 }
 
 func Execute() {
+	rootCmd := NewRootCommand()
 	if err := rootCmd.Execute(); err != nil {
 		if merr, ok := err.(*multierror.Error); ok {
 			for _, e := range merr.Errors {

cmd/root_test.go

@@ -0,0 +1,189 @@
+package cmd
+
+import (
+	"fmt"
+	"log"
+	"net"
+	"os"
+	"testing"
+
+	"github.com/miekg/dns"
+	"github.com/stretchr/testify/assert"
+)
+
+const (
+	DNSServerPort = 53535
+)
+
+func TestMain(m *testing.M) {
+	go startDNSServer()
+
+	code := m.Run()
+	os.Exit(code)
+}
+
+func startDNSServer() {
+	dns.HandleFunc(".", dnsHandler)
+
+	go func() {
+		err := dns.ListenAndServe(fmt.Sprintf(":%d", DNSServerPort), "udp", nil)
+		if err != nil {
+			log.Fatalf("Failed to start DNS server: %v", err)
+		}
+	}()
+}
+
+func dnsHandler(w dns.ResponseWriter, r *dns.Msg) {
+	msg := dns.Msg{}
+	msg.SetReply(r)
+
+	// Simulate an A record response for "example.com"
+	if len(r.Question) > 0 {
+		q := r.Question[0]
+		if q.Name == "example.com." && q.Qtype == dns.TypeA {
+			// Example A record response
+			a := &dns.A{
+				Hdr: dns.RR_Header{
+					Name:   "example.com.",
+					Rrtype: dns.TypeA,
+					Class:  dns.ClassINET,
+					Ttl:    60,
+				},
+				A: net.ParseIP("93.184.216.34"),
+			}
+			msg.Answer = append(msg.Answer, a)
+		}
+		// Simulate a CNAME record response for "example.com"
+		if q.Name == "example.com." && q.Qtype == dns.TypeCNAME {
+			// Example CNAME record response
+			cname := &dns.CNAME{
+				Hdr: dns.RR_Header{
+					Name:   "example.com.",
+					Rrtype: dns.TypeCNAME,
+					Class:  dns.ClassINET,
+					Ttl:    60,
+				},
+				Target: "example.org.",
+			}
+			msg.Answer = append(msg.Answer, cname)
+		}
+	}
+
+	_ = w.WriteMsg(&msg)
+}
+
+func Test_Cmd(t *testing.T) {
+	t.Setenv("NO_COLOR", "1") // Disable color codes for easier testing
+
+	rootCmd := NewRootCommand()
+	rootCmd.SetArgs([]string{"example.com", "--server", fmt.Sprintf("127.0.0.1:%d", DNSServerPort)})
+
+	err := rootCmd.Execute()
+
+	assert.NoError(t, err)
+}
+
+func Test_Cmd_Error(t *testing.T) {
+	t.Setenv("NO_COLOR", "1") // Disable color codes for easier testing
+
+	rootCmd := NewRootCommand()
+	rootCmd.SetArgs([]string{"--server", fmt.Sprintf("127.0.0.1:%d", DNSServerPort)})
+
+	err := rootCmd.Execute()
+
+	assert.Error(t, err)
+	assert.Equal(t, "error: domain name is required", err.Error())
+}
+
+func Test_Cmd_JSON(t *testing.T) {
+	t.Setenv("NO_COLOR", "1") // Disable color codes for easier testing
+
+	rootCmd := NewRootCommand()
+	rootCmd.SetArgs([]string{"example.com", "--json", "--server", fmt.Sprintf("127.0.0.1:%d", DNSServerPort)})
+
+	err := rootCmd.Execute()
+
+	assert.NoError(t, err)
+}
+
+func Test_Cmd_QueryType(t *testing.T) {
+	t.Setenv("NO_COLOR", "1") // Disable color codes for easier testing
+
+	rootCmd := NewRootCommand()
+	rootCmd.SetArgs([]string{"example.com", "--server", fmt.Sprintf("127.0.0.1:%d", DNSServerPort), "--query-type", "A"})
+
+	err := rootCmd.Execute()
+
+	assert.NoError(t, err)
+}
+
+func Test_Cmd_Debug(t *testing.T) {
+	t.Setenv("NO_COLOR", "1") // Disable color codes for easier testing
+
+	rootCmd := NewRootCommand()
+	rootCmd.SetArgs([]string{"example.com", "--debug", "--server", fmt.Sprintf("127.0.0.1:%d", DNSServerPort), "--query-type", "A"})
+
+	err := rootCmd.Execute()
+
+	assert.NoError(t, err)
+}
+
+func Test_Cmd_LogFile(t *testing.T) {
+	t.Setenv("NO_COLOR", "1") // Disable color codes for easier testing
+
+	file, err := os.CreateTemp(t.TempDir(), "zns")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer os.Remove(file.Name())
+
+	t.Setenv("ZNS_LOG_FILE", file.Name())
+
+	rootCmd := NewRootCommand()
+	rootCmd.SetArgs([]string{"example.com", "--server", fmt.Sprintf("127.0.0.1:%d", DNSServerPort)})
+
+	err = rootCmd.Execute()
+	assert.NoError(t, err)
+
+	assert.FileExists(t, file.Name())
+
+	logFile, err := os.ReadFile(file.Name())
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	assert.Contains(t, string(logFile), "A       example.com.   01m00s   93.184.216.34")
+	assert.Contains(t, string(logFile), "CNAME   example.com.   01m00s   example.org.")
+}
+
+func Test_Cmd_LogFile_Debug(t *testing.T) {
+	t.Setenv("NO_COLOR", "1") // Disable color codes for easier testing
+
+	file, err := os.CreateTemp(t.TempDir(), "zns")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer os.Remove(file.Name())
+
+	t.Setenv("ZNS_LOG_FILE", file.Name())
+
+	rootCmd := NewRootCommand()
+	rootCmd.SetArgs([]string{"example.com", "--debug", "--server", fmt.Sprintf("127.0.0.1:%d", DNSServerPort)})
+
+	err = rootCmd.Execute()
+	assert.NoError(t, err)
+
+	assert.FileExists(t, file.Name())
+
+	logFile, err := os.ReadFile(file.Name())
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	assert.Contains(t, string(logFile), "Querying DNS server: @domain=example.com server=127.0.0.1:53535 domain=example.com qtype=A")
+	assert.Contains(t, string(logFile), "Querying DNS server: @domain=example.com server=127.0.0.1:53535 domain=example.com qtype=CNAME")
+	assert.Contains(t, string(logFile), "Received DNS response: @domain=example.com server=127.0.0.1:53535 domain=example.com qtype=A rcode=NOERROR")
+	assert.Contains(t, string(logFile), "Received DNS response: @domain=example.com server=127.0.0.1:53535 domain=example.com qtype=CNAME rcode=NOERROR")
+	assert.Contains(t, string(logFile), "A       |example.com.   |01m00s   |93.184.216.34")
+	assert.Contains(t, string(logFile), "CNAME   |example.com.   |01m00s   |example.org.")
+}