coredns/plugin/file/tree/print_test.go

package tree

import (
	"net"
	"os"
	"strings"
	"testing"

	"github.com/miekg/dns"
)

func TestPrint(t *testing.T) {
	rr1 := dns.A{
		Hdr: dns.RR_Header{
			Name:     dns.Fqdn("server1.example.com"),
			Rrtype:   1,
			Class:    1,
			Ttl:      3600,
			Rdlength: 0,
		},
		A: net.IPv4(10, 0, 1, 1),
	}
	rr2 := dns.A{
		Hdr: dns.RR_Header{
			Name:     dns.Fqdn("server2.example.com"),
			Rrtype:   1,
			Class:    1,
			Ttl:      3600,
			Rdlength: 0,
		},
		A: net.IPv4(10, 0, 1, 2),
	}
	rr3 := dns.A{
		Hdr: dns.RR_Header{
			Name:     dns.Fqdn("server3.example.com"),
			Rrtype:   1,
			Class:    1,
			Ttl:      3600,
			Rdlength: 0,
		},
		A: net.IPv4(10, 0, 1, 3),
	}
	rr4 := dns.A{
		Hdr: dns.RR_Header{
			Name:     dns.Fqdn("server4.example.com"),
			Rrtype:   1,
			Class:    1,
			Ttl:      3600,
			Rdlength: 0,
		},
		A: net.IPv4(10, 0, 1, 4),
	}
	tree := Tree{
		Root:  nil,
		Count: 0,
	}
	tree.Insert(&rr1)
	tree.Insert(&rr2)
	tree.Insert(&rr3)
	tree.Insert(&rr4)

	/**
	build a LLRB tree, the height of the tree is 3, look like:

				  server2.example.com.
					/             \
		server1.example.com.   server4.example.com.
			   /
	 server3.example.com.

	*/

	f, err := os.CreateTemp(t.TempDir(), "print_test_tmp")
	if err != nil {
		t.Fatal(err)
	}
	// Redirect Print output to a temp file, then restore stdout and close
	// the handle so TempDir cleanup can remove it on Windows.
	stdout := os.Stdout
	os.Stdout = f
	t.Cleanup(func() {
		os.Stdout = stdout
		f.Close()
	})

	tree.Print()
	/**
	  server2.example.com.
	  server1.example.com. server4.example.com.
	  server3.example.com.
	*/

	if _, err := f.Seek(0, 0); err != nil {
		t.Fatal(err)
	}
	buf := make([]byte, 256)
	if _, err := f.Read(buf); err != nil {
		t.Fatal(err)
	}
	height := strings.Count(string(buf), ". \n")
	//Compare the height of the print with the actual height of the tree
	if height != 3 {
		t.Fatal("The number of rows is inconsistent with the actual number of rows in the tree itself.")
	}
}