Test that Msg.RemoveRRset() works as intended

update_test.go

@@ -1,6 +1,9 @@
 package dns
 
-import "testing"
+import (
+	"bytes"
+	"testing"
+)
 
 func TestDynamicUpdateParsing(t *testing.T) {
 	prefix := "example.com. IN "
@@ -52,3 +55,51 @@ func TestDynamicUpdateZeroRdataUnpack(t *testing.T) {
 		}
 	}
 }
+
+func TestRemoveRRset(t *testing.T) {
+	// Should add a zero data RR in Class ANY with a TTL of 0
+	// for each set mentioned in the RRs provided to it.
+	rr, err := NewRR(". 100 IN A 127.0.0.1")
+	if err != nil {
+		t.Fatalf("Error constructing RR: %v", err)
+	}
+	m := new(Msg)
+	m.Ns = []RR{&RR_Header{Name: ".", Rrtype: TypeA, Class: ClassANY, Ttl: 0, Rdlength: 0}}
+	expectstr := m.String()
+	expect, err := m.Pack()
+	if err != nil {
+		t.Fatalf("Error packing expected msg: %v", err)
+	}
+
+	m.Ns = nil
+	m.RemoveRRset([]RR{rr})
+	actual, err := m.Pack()
+	if err != nil {
+		t.Fatalf("Error packing actual msg: %v", err)
+	}
+	if !bytes.Equal(actual, expect) {
+		tmp := new(Msg)
+		if err := tmp.Unpack(actual); err != nil {
+			t.Fatalf("Error unpacking actual msg: %v", err)
+		}
+		t.Logf("Expected msg:\n%s", expectstr)
+		t.Logf("Actual msg:\n%v", tmp)
+		t.Fail()
+	}
+
+	m.Ns = nil
+	m.RemoveRRset([]RR{rr, rr})
+	actual, err = m.Pack()
+	if err != nil {
+		t.Fatalf("Error packing actual msg: %v", err)
+	}
+	if !bytes.Equal(actual, expect) {
+		tmp := new(Msg)
+		if err := tmp.Unpack(actual); err != nil {
+			t.Fatalf("Error unpacking actual msg: %v", err)
+		}
+		t.Logf("Expected msg:\n%v", expectstr)
+		t.Logf("Actual msg:\n%v", tmp)
+		t.Fail()
+	}
+}