1 files changed, 31 insertions, 9 deletions
diff --git a/request/request_test.go b/request/request_test.go
index 49b825627..e6fcf6935 100644
--- a/request/request_test.go
+++ b/request/request_test.go
@@ -61,7 +61,7 @@ func TestRequestMalformed(t *testing.T) {
 	}
 }
 
-func TestRequestScrub(t *testing.T) {
+func TestRequestScrubAnswer(t *testing.T) {
 	m := new(dns.Msg)
 	m.SetQuestion("large.example.com.", dns.TypeSRV)
 	req := Request{W: &test.ResponseWriter{}, Req: m}
@@ -69,24 +69,46 @@ func TestRequestScrub(t *testing.T) {
 	reply := new(dns.Msg)
 	reply.SetReply(m)
 	for i := 1; i < 200; i++ {
-		reply.Answer = append(reply.Answer, test.SRV(fmt.Sprintf(
-			"large.example.com. 10 IN SRV 0 0 80 10-0-0-%d.default.pod.k8s.example.com.",
-			i,
-		)))
+		reply.Answer = append(reply.Answer, test.SRV(
+			fmt.Sprintf("large.example.com. 10 IN SRV 0 0 80 10-0-0-%d.default.pod.k8s.example.com.", i)))
 	}
 
-	msg, got := req.Scrub(reply)
-	if want := ScrubDone; want != got {
+	_, got := req.Scrub(reply)
+	if want := ScrubAnswer; want != got {
 		t.Errorf("want scrub result %d, got %d", want, got)
 	}
-	if want, got := req.Size(), msg.Len(); want < got {
+	if want, got := req.Size(), reply.Len(); want < got {
 		t.Errorf("want scrub to reduce message length below %d bytes, got %d bytes", want, got)
 	}
-	if !msg.Truncated {
+	if !reply.Truncated {
 		t.Errorf("want scrub to set truncated bit")
 	}
 }
 
+func TestRequestScrubExtra(t *testing.T) {
+	m := new(dns.Msg)
+	m.SetQuestion("large.example.com.", dns.TypeSRV)
+	req := Request{W: &test.ResponseWriter{}, Req: m}
+
+	reply := new(dns.Msg)
+	reply.SetReply(m)
+	for i := 1; i < 200; i++ {
+		reply.Extra = append(reply.Extra, test.SRV(
+			fmt.Sprintf("large.example.com. 10 IN SRV 0 0 80 10-0-0-%d.default.pod.k8s.example.com.", i)))
+	}
+
+	_, got := req.Scrub(reply)
+	if want := ScrubExtra; want != got {
+		t.Errorf("want scrub result %d, got %d", want, got)
+	}
+	if want, got := req.Size(), reply.Len(); want < got {
+		t.Errorf("want scrub to reduce message length below %d bytes, got %d bytes", want, got)
+	}
+	if reply.Truncated {
+		t.Errorf("want scrub to not set truncated bit")
+	}
+}
+
 func BenchmarkRequestDo(b *testing.B) {
 	st := testRequest()