From 07b82d399371f688cc8c22631088fef18128097d Mon Sep 17 00:00:00 2001
From: Steven Lee <steven.lee@kynrai.com>
Date: Mon, 10 Jun 2024 15:35:50 +0100
Subject: [PATCH] RF: use testify in all tests

---
 handlers/cookies_test.go    | 45 ++++++++++---------------------------
 handlers/dns_server_test.go | 12 ++++------
 handlers/dns_test.go        |  6 ++---
 handlers/dnssec_test.go     |  6 ++---
 handlers/firewall_test.go   |  6 ++---
 handlers/redirects_test.go  |  1 +
 6 files changed, 26 insertions(+), 50 deletions(-)

diff --git a/handlers/cookies_test.go b/handlers/cookies_test.go
index db343b6..7b460ae 100644
--- a/handlers/cookies_test.go
+++ b/handlers/cookies_test.go
@@ -3,45 +3,24 @@ package handlers
 import (
 	"net/http"
 	"net/http/httptest"
-	"strings"
 	"testing"
+
+	"github.com/stretchr/testify/assert"
 )
 
 func TestHandlerCookies(t *testing.T) {
 	t.Parallel()
-	testCases := []struct {
-		name          string
-		url           string
-		expectedCode  int
-		expectedError string
-	}{
-		{
-			name:         "Missing URL",
-			url:          "",
-			expectedCode: http.StatusBadRequest,
-		},
-		{
-			name:         "Invalid URL",
-			url:          "invalid_url",
-			expectedCode: http.StatusInternalServerError,
-		},
-	}
 
-	for _, tc := range testCases {
-		tc := tc
-		t.Run(tc.name, func(t *testing.T) {
-			t.Parallel()
-			req := httptest.NewRequest("GET", "/cookies?url="+tc.url, nil)
-			rec := httptest.NewRecorder()
-			HandleCookies().ServeHTTP(rec, req)
+	t.Run("missing URL parameter", func(t *testing.T) {
+		t.Parallel()
+		req := httptest.NewRequest(http.MethodGet, "/blocklists", nil)
+		rec := httptest.NewRecorder()
+
+		HandleCookies().ServeHTTP(rec, req)
 
-			if rec.Code != tc.expectedCode {
-				t.Errorf("Expected status code %d, got %d", tc.expectedCode, rec.Code)
-			}
+		assert.Equal(t, http.StatusBadRequest, rec.Code)
+		assert.JSONEq(t, `{"error": "missing URL parameter"}`, rec.Body.String())
+	})
 
-			if tc.expectedError != "" && !strings.Contains(rec.Body.String(), tc.expectedError) {
-				t.Errorf("Expected error message '%s' not found in response body", tc.expectedError)
-			}
-		})
-	}
+	// TODO: Add tests for cookie logic
 }
diff --git a/handlers/dns_server_test.go b/handlers/dns_server_test.go
index 093ee5a..0245abc 100644
--- a/handlers/dns_server_test.go
+++ b/handlers/dns_server_test.go
@@ -3,8 +3,9 @@ package handlers
 import (
 	"net/http"
 	"net/http/httptest"
-	"strings"
 	"testing"
+
+	"github.com/stretchr/testify/assert"
 )
 
 func TestHandleDNSServer(t *testing.T) {
@@ -38,13 +39,8 @@ func TestHandleDNSServer(t *testing.T) {
 
 			HandleDNSServer().ServeHTTP(rec, req)
 
-			if rec.Code != tc.expectedCode {
-				t.Errorf("Expected status code %d, got %d", tc.expectedCode, rec.Code)
-			}
-
-			if strings.TrimSpace(rec.Body.String()) != tc.expectedBody {
-				t.Errorf("Expected body '%s', got '%s'", tc.expectedBody, strings.TrimSpace(rec.Body.String()))
-			}
+			assert.Equal(t, tc.expectedCode, rec.Code)
+			assert.JSONEq(t, tc.expectedBody, rec.Body.String())
 		})
 	}
 }
diff --git a/handlers/dns_test.go b/handlers/dns_test.go
index 29f5308..a77520b 100644
--- a/handlers/dns_test.go
+++ b/handlers/dns_test.go
@@ -4,6 +4,8 @@ import (
 	"net/http"
 	"net/http/httptest"
 	"testing"
+
+	"github.com/stretchr/testify/assert"
 )
 
 func TestHandleDNS(t *testing.T) {
@@ -34,9 +36,7 @@ func TestHandleDNS(t *testing.T) {
 
 			HandleDNS().ServeHTTP(rec, req)
 
-			if rec.Code != tc.expectedCode {
-				t.Errorf("Expected status code %d, got %d", tc.expectedCode, rec.Code)
-			}
+			assert.Equal(t, tc.expectedCode, rec.Code)
 		})
 	}
 }
diff --git a/handlers/dnssec_test.go b/handlers/dnssec_test.go
index 15e0daa..2040add 100644
--- a/handlers/dnssec_test.go
+++ b/handlers/dnssec_test.go
@@ -4,6 +4,8 @@ import (
 	"net/http"
 	"net/http/httptest"
 	"testing"
+
+	"github.com/stretchr/testify/assert"
 )
 
 func TestHandleDnsSec(t *testing.T) {
@@ -34,9 +36,7 @@ func TestHandleDnsSec(t *testing.T) {
 
 			HandleDnsSec().ServeHTTP(rec, req)
 
-			if rec.Code != tc.expectedCode {
-				t.Errorf("Expected status code %d, got %d", tc.expectedCode, rec.Code)
-			}
+			assert.Equal(t, tc.expectedCode, rec.Code)
 		})
 	}
 }
diff --git a/handlers/firewall_test.go b/handlers/firewall_test.go
index 44698e2..c2aa23c 100644
--- a/handlers/firewall_test.go
+++ b/handlers/firewall_test.go
@@ -4,6 +4,8 @@ import (
 	"net/http"
 	"net/http/httptest"
 	"testing"
+
+	"github.com/stretchr/testify/assert"
 )
 
 func TestHandleFirewall(t *testing.T) {
@@ -34,9 +36,7 @@ func TestHandleFirewall(t *testing.T) {
 
 			HandleFirewall().ServeHTTP(rec, req)
 
-			if rec.Code != tc.expectedCode {
-				t.Errorf("Expected status code %d, got %d", tc.expectedCode, rec.Code)
-			}
+			assert.Equal(t, tc.expectedCode, rec.Code)
 		})
 	}
 }
diff --git a/handlers/redirects_test.go b/handlers/redirects_test.go
index a2af671..770f390 100644
--- a/handlers/redirects_test.go
+++ b/handlers/redirects_test.go
@@ -50,6 +50,7 @@ func TestHandleGetRedirects(t *testing.T) {
 			var response map[string]interface{}
 			err := json.Unmarshal(rec.Body.Bytes(), &response)
 			assert.NoError(t, err)
+			// TODO: break this out of table drive tests, should not use name as part of logic
 			if tc.name == "Invalid URL" {
 				assert.Contains(t, response["error"], tc.expectedBody["error"])
 			} else {