dns_test: Increase coverage for dns.go

Signed-off-by: Simarpreet Singh <[email protected]>

Author: simar7

dns.go

@@ -51,23 +51,23 @@ func serve(w dns.ResponseWriter, r *dns.Msg) {
 	w.WriteMsg(m)
 }
 
-func newDNSServer(addr string, port int, conntype string) *dns.Server {
+func newDNSServerHandler(addr string, port int, conntype string) *dns.Server {
 	return &dns.Server{
 		Addr: fmt.Sprintf("%s:%d", addr, port),
 		Net:  conntype,
 	}
 }
 
-func newDefaultDNSServer() *dns.Server {
-	return &dns.Server{
-		Addr: fmt.Sprintf("%s:%d", DNS_SERVER_ADDR, DNS_SERVER_PORT),
-		Net:  "udp",
+func newDefaultDNSServerHandler() *dnsServerHandler {
+	return &dnsServerHandler{
+		ds: newDNSServerHandler(DNS_SERVER_ADDR, DNS_SERVER_PORT, "udp"),
 	}
 }
 
 func (h *dnsServerHandler) RouteDNS() error {
 	rl := NewLogger(log.New().Writer())
 
+	// TODO: Add more domains to the serving list
 	dns.HandleFunc("com.", serve)
 	err := h.ds.ListenAndServe()
 	if err != nil {

dns_test.go

@@ -0,0 +1,73 @@
+package main
+
+import (
+	"net"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/miekg/dns"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func RunLocalUDPServer(laddr string) (*dns.Server, string, error) {
+	server, l, _, err := RunLocalUDPServerWithFinChan(laddr)
+
+	return server, l, err
+}
+
+func RunLocalUDPServerWithFinChan(laddr string) (*dns.Server, string, chan struct{}, error) {
+	pc, err := net.ListenPacket("udp", laddr)
+	if err != nil {
+		return nil, "", nil, err
+	}
+	server := &dns.Server{PacketConn: pc, ReadTimeout: time.Hour, WriteTimeout: time.Hour}
+
+	waitLock := sync.Mutex{}
+	waitLock.Lock()
+	server.NotifyStartedFunc = waitLock.Unlock
+
+	fin := make(chan struct{}, 0)
+
+	go func() {
+		server.ActivateAndServe()
+		close(fin)
+		pc.Close()
+	}()
+
+	waitLock.Lock()
+	return server, pc.LocalAddr().String(), fin, nil
+}
+
+func TestDNS_OK(t *testing.T) {
+	h := newDefaultDNSServerHandler()
+	dns.HandleFunc("test.", serve)
+	defer dns.HandleRemove("test.")
+
+	err := h.ds.ListenAndServe()
+	if err != nil {
+		assert.Error(t, err)
+	}
+
+	m := new(dns.Msg)
+	m.SetQuestion("go.test.", dns.TypeA)
+
+	s, addrstr, err := RunLocalUDPServer("127.0.0.1:0")
+	if err != nil {
+		t.Fatalf("unable to run test server: %v", err)
+	}
+	defer s.Shutdown()
+
+	cn, err := dns.Dial("udp", addrstr)
+	if err != nil {
+		t.Errorf("failed to dial %s: %v", addrstr, err)
+	}
+	err = cn.WriteMsg(m)
+	assert.NoError(t, err)
+
+	r, err := cn.ReadMsg()
+	assert.NoError(t, err)
+	assert.NotNil(t, r)
+	assert.NotEqual(t, r, dns.RcodeSuccess)
+}

main.go

@@ -27,9 +27,7 @@ func main() {
 	rl := NewLogger(log.New().Writer())
 
 	dhcpHandler := newDHCPServer()
-	dnsHandler := dnsServerHandler{
-		ds: newDefaultDNSServer(),
-	}
+	dnsHandler := newDefaultDNSServerHandler()
 
 	go func() {
 		wg.Add(1)