plugin/file: Fix print tree error (#4962)

Automatically submitted.

2 files changed, 104 insertions, 2 deletions

diff --git a/plugin/file/tree/print.go b/plugin/file/tree/print.go
index bd86ef690..b2df70e1b 100644
--- a/plugin/file/tree/print.go
+++ b/plugin/file/tree/print.go
@@ -29,9 +29,9 @@ func (n *Node) print() {
 		}
 		if nodesInCurrentLevel == 0 {
 			fmt.Println()
+			nodesInCurrentLevel = nodesInNextLevel
+			nodesInNextLevel = 0
 		}
-		nodesInCurrentLevel = nodesInNextLevel
-		nodesInNextLevel = 0
 	}
 	fmt.Println()
 }
diff --git a/plugin/file/tree/print_test.go b/plugin/file/tree/print_test.go
new file mode 100644
index 000000000..2ab552769
--- /dev/null
+++ b/plugin/file/tree/print_test.go
@@ -0,0 +1,102 @@
+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.Create("tmp")
+	if err != nil {
+		t.Error(err)
+	}
+	//Redirect the printed results to a tmp file for later comparison
+	os.Stdout = f
+
+	tree.Print()
+	/**
+	  server2.example.com.
+	  server1.example.com. server4.example.com.
+	  server3.example.com.
+	*/
+
+	buf := make([]byte, 256)
+	f.Seek(0, 0)
+	_, er := f.Read(buf)
+	if er != nil {
+		t.Error(err)
+	}
+	height := strings.Count(string(buf), ". \n")
+	//Compare the height of the print with the actual height of the tree
+	if height != 3 {
+		f.Close()
+		os.Remove("tmp")
+		t.Fatal("The number of rows is inconsistent with the actual number of rows in the tree itself.")
+	}
+	f.Close()
+	os.Remove("tmp")
+}