From 6e39d1e5f47dcd67258da2b950b1126556709662 Mon Sep 17 00:00:00 2001
From: Sam <sam@samtherapy.net>
Date: Wed, 28 Sep 2022 14:49:52 +0000
Subject: [PATCH] test: Add retries to all tests

Add the retry loop from main in the event of lone failures.

Hopefully this fixes test flakiness at the cost of taking longer
---
 pkg/query/print_test.go        | 16 +++++++++++++---
 pkg/resolvers/DNSCrypt_test.go | 14 +++++++++++++-
 pkg/resolvers/HTTPS_test.go    | 18 +++++++++++-------
 pkg/resolvers/QUIC_test.go     | 20 ++++++++++++--------
 pkg/resolvers/general_test.go  | 14 +++++++++++++-
 5 files changed, 62 insertions(+), 20 deletions(-)

diff --git a/pkg/query/print_test.go b/pkg/query/print_test.go
index a758f8a..4b3c8b8 100644
--- a/pkg/query/print_test.go
+++ b/pkg/query/print_test.go
@@ -187,16 +187,26 @@ func TestRealPrint(t *testing.T) {
 
 		t.Run("", func(t *testing.T) {
 			t.Parallel()
-			resp, err := query.CreateQuery(test)
+
+			var (
+				res util.Response
+				err error
+			)
+			for i := 0; i <= test.Request.Retries; i++ {
+				res, err = query.CreateQuery(test)
+				if err == nil {
+					break
+				}
+			}
 			assert.NilError(t, err)
 
 			if test.JSON || test.XML || test.YAML {
 				str := ""
-				str, err = query.PrintSpecial(resp, test)
+				str, err = query.PrintSpecial(res, test)
 				assert.NilError(t, err)
 				assert.Assert(t, str != "")
 			}
-			str, err := query.ToString(resp, test)
+			str, err := query.ToString(res, test)
 			assert.NilError(t, err)
 			assert.Assert(t, str != "")
 		})
diff --git a/pkg/resolvers/DNSCrypt_test.go b/pkg/resolvers/DNSCrypt_test.go
index 98f5e2f..df1a731 100644
--- a/pkg/resolvers/DNSCrypt_test.go
+++ b/pkg/resolvers/DNSCrypt_test.go
@@ -3,10 +3,12 @@
 package resolvers_test
 
 import (
+	"errors"
 	"testing"
 
 	"git.froth.zone/sam/awl/pkg/query"
 	"git.froth.zone/sam/awl/pkg/util"
+	"github.com/ameshkov/dnscrypt/v2"
 	"github.com/miekg/dns"
 	"gotest.tools/v3/assert"
 )
@@ -69,7 +71,17 @@ func TestDNSCrypt(t *testing.T) {
 		t.Run(test.name, func(t *testing.T) {
 			t.Parallel()
 
-			res, err := query.CreateQuery(test.opts)
+			var (
+				res util.Response
+				err error
+			)
+			for i := 0; i <= test.opts.Request.Retries; i++ {
+				res, err = query.CreateQuery(test.opts)
+				if err == nil || errors.Is(err, dnscrypt.ErrInvalidDNSStamp) {
+					break
+				}
+			}
+
 			if err == nil {
 				assert.Assert(t, res != util.Response{})
 			} else {
diff --git a/pkg/resolvers/HTTPS_test.go b/pkg/resolvers/HTTPS_test.go
index 2c3742a..5505c58 100644
--- a/pkg/resolvers/HTTPS_test.go
+++ b/pkg/resolvers/HTTPS_test.go
@@ -6,6 +6,7 @@ import (
 	"errors"
 	"testing"
 
+	"git.froth.zone/sam/awl/pkg/query"
 	"git.froth.zone/sam/awl/pkg/resolvers"
 	"git.froth.zone/sam/awl/pkg/util"
 	"github.com/miekg/dns"
@@ -76,13 +77,16 @@ func TestHTTPS(t *testing.T) {
 		t.Run(test.name, func(t *testing.T) {
 			t.Parallel()
 
-			resolver, err := resolvers.LoadResolver(test.opts)
-			assert.NilError(t, err)
-
-			msg := new(dns.Msg)
-			msg.SetQuestion(test.opts.Request.Name, test.opts.Request.Type)
-			// msg = msg.SetQuestion(testCase.Name, testCase.Type)
-			res, err := resolver.LookUp(msg)
+			var (
+				res util.Response
+				err error
+			)
+			for i := 0; i <= test.opts.Request.Retries; i++ {
+				res, err = query.CreateQuery(test.opts)
+				if err == nil || errors.Is(err, &resolvers.ErrHTTPStatus{}) {
+					break
+				}
+			}
 
 			if err == nil {
 				assert.NilError(t, err)
diff --git a/pkg/resolvers/QUIC_test.go b/pkg/resolvers/QUIC_test.go
index 0691558..029df11 100644
--- a/pkg/resolvers/QUIC_test.go
+++ b/pkg/resolvers/QUIC_test.go
@@ -6,7 +6,7 @@ import (
 	"testing"
 	"time"
 
-	"git.froth.zone/sam/awl/pkg/resolvers"
+	"git.froth.zone/sam/awl/pkg/query"
 	"git.froth.zone/sam/awl/pkg/util"
 	"github.com/miekg/dns"
 	"gotest.tools/v3/assert"
@@ -74,6 +74,7 @@ func TestQuic(t *testing.T) {
 					Type:    dns.TypeA,
 					Name:    "git.froth.zone",
 					Timeout: 10 * time.Millisecond,
+					Retries: 0,
 				},
 			},
 		},
@@ -85,13 +86,16 @@ func TestQuic(t *testing.T) {
 		t.Run(test.name, func(t *testing.T) {
 			t.Parallel()
 
-			resolver, err := resolvers.LoadResolver(test.opts)
-			assert.NilError(t, err)
-
-			msg := new(dns.Msg)
-			msg.SetQuestion(test.opts.Request.Name, test.opts.Request.Type)
-
-			res, err := resolver.LookUp(msg)
+			var (
+				res util.Response
+				err error
+			)
+			for i := 0; i <= test.opts.Request.Retries; i++ {
+				res, err = query.CreateQuery(test.opts)
+				if err == nil {
+					break
+				}
+			}
 
 			if err == nil {
 				assert.NilError(t, err)
diff --git a/pkg/resolvers/general_test.go b/pkg/resolvers/general_test.go
index 6e7774d..af4b16a 100644
--- a/pkg/resolvers/general_test.go
+++ b/pkg/resolvers/general_test.go
@@ -3,12 +3,14 @@
 package resolvers_test
 
 import (
+	"errors"
 	"os"
 	"testing"
 	"time"
 
 	"git.froth.zone/sam/awl/pkg/query"
 	"git.froth.zone/sam/awl/pkg/util"
+	"github.com/ameshkov/dnscrypt/v2"
 	"github.com/miekg/dns"
 	"gotest.tools/v3/assert"
 )
@@ -116,7 +118,17 @@ func TestResolve(t *testing.T) {
 		t.Run(test.name, func(t *testing.T) {
 			t.Parallel()
 
-			res, err := query.CreateQuery(test.opts)
+			var (
+				res util.Response
+				err error
+			)
+			for i := 0; i <= test.opts.Request.Retries; i++ {
+				res, err = query.CreateQuery(test.opts)
+				if err == nil || errors.Is(err, dnscrypt.ErrInvalidDNSStamp) {
+					break
+				}
+			}
+
 			if err == nil {
 				assert.NilError(t, err)
 				assert.Assert(t, res != util.Response{})