test: Fix flaky DNS tests that check for valid IP (#968)

TwiN · Jan 19, 2025 · 69dbe4f · 69dbe4f
1 parent 9157b5b
commit 69dbe4f
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 17 deletions.
diff --git a/client/client_test.go b/client/client_test.go
@@ -5,6 +5,7 @@ import (
 	"crypto/tls"
 	"io"
 	"net/http"
+	"net/netip"
 	"testing"
 	"time"
 
@@ -356,7 +357,7 @@ func TestTlsRenegotiation(t *testing.T) {
 }
 
 func TestQueryDNS(t *testing.T) {
-	tests := []struct {
+	scenarios := []struct {
 		name            string
 		inputDNS        dns.Config
 		inputURL        string
@@ -372,7 +373,7 @@ func TestQueryDNS(t *testing.T) {
 			},
 			inputURL:        "8.8.8.8",
 			expectedDNSCode: "NOERROR",
-			expectedBody:    "93.184.215.14",
+			expectedBody:    "__IPV4__",
 		},
 		{
 			name: "test Config with type AAAA",
@@ -382,7 +383,7 @@ func TestQueryDNS(t *testing.T) {
 			},
 			inputURL:        "8.8.8.8",
 			expectedDNSCode: "NOERROR",
-			expectedBody:    "2606:2800:21f:cb07:6820:80da:af6b:8b2c",
+			expectedBody:    "__IPV6__",
 		},
 		{
 			name: "test Config with type CNAME",
@@ -434,27 +435,42 @@ func TestQueryDNS(t *testing.T) {
 			isErrExpected: true,
 		},
 	}
-
-	for _, test := range tests {
-		t.Run(test.name, func(t *testing.T) {
-			_, dnsRCode, body, err := QueryDNS(test.inputDNS.QueryType, test.inputDNS.QueryName, test.inputURL)
-			if test.isErrExpected && err == nil {
+	for _, scenario := range scenarios {
+		t.Run(scenario.name, func(t *testing.T) {
+			_, dnsRCode, body, err := QueryDNS(scenario.inputDNS.QueryType, scenario.inputDNS.QueryName, scenario.inputURL)
+			if scenario.isErrExpected && err == nil {
 				t.Errorf("there should be an error")
 			}
-			if dnsRCode != test.expectedDNSCode {
-				t.Errorf("expected DNSRCode to be %s, got %s", test.expectedDNSCode, dnsRCode)
+			if dnsRCode != scenario.expectedDNSCode {
+				t.Errorf("expected DNSRCode to be %s, got %s", scenario.expectedDNSCode, dnsRCode)
 			}
-			if test.inputDNS.QueryType == "NS" {
+			if scenario.inputDNS.QueryType == "NS" {
 				// Because there are often multiple nameservers backing a single domain, we'll only look at the suffix
-				if !pattern.Match(test.expectedBody, string(body)) {
-					t.Errorf("got %s, expected result %s,", string(body), test.expectedBody)
+				if !pattern.Match(scenario.expectedBody, string(body)) {
+					t.Errorf("got %s, expected result %s,", string(body), scenario.expectedBody)
 				}
 			} else {
-				if string(body) != test.expectedBody {
-					t.Errorf("got %s, expected result %s,", string(body), test.expectedBody)
+				if string(body) != scenario.expectedBody {
+					// little hack to validate arbitrary ipv4/ipv6
+					switch scenario.expectedBody {
+					case "__IPV4__":
+						if addr, err := netip.ParseAddr(string(body)); err != nil {
+							t.Errorf("got %s, expected result %s", string(body), scenario.expectedBody)
+						} else if !addr.Is4() {
+							t.Errorf("got %s, expected valid IPv4", string(body))
+						}
+					case "__IPV6__":
+						if addr, err := netip.ParseAddr(string(body)); err != nil {
+							t.Errorf("got %s, expected result %s", string(body), scenario.expectedBody)
+						} else if !addr.Is6() {
+							t.Errorf("got %s, expected valid IPv6", string(body))
+						}
+					default:
+						t.Errorf("got %s, expected result %s", string(body), scenario.expectedBody)
+					}
 				}
 			}
 		})
-		time.Sleep(5 * time.Millisecond)
+		time.Sleep(10 * time.Millisecond)
 	}
 }
diff --git a/config/endpoint/endpoint_test.go b/config/endpoint/endpoint_test.go
@@ -761,7 +761,7 @@ func TestIntegrationEvaluateHealthWithErrorAndHideURL(t *testing.T) {
 
 func TestIntegrationEvaluateHealthForDNS(t *testing.T) {
 	conditionSuccess := Condition("[DNS_RCODE] == NOERROR")
-	conditionBody := Condition("[BODY] == 93.184.215.14")
+	conditionBody := Condition("[BODY] == pat(*.*.*.*)")
 	endpoint := Endpoint{
 		Name: "example",
 		URL:  "8.8.8.8",