From d20d0bb33746abd2be92b5f17a8e226183c6e0e7 Mon Sep 17 00:00:00 2001
From: petruki <31597636+petruki@users.noreply.github.com>
Date: Wed, 13 May 2026 19:41:05 -0700
Subject: [PATCH] chore(tests): fixed formatting issues, simplified helpers

---
 snapshot_test.go | 144 ++++++++++++++++++++++++++---------------------
 1 file changed, 80 insertions(+), 64 deletions(-)

diff --git a/snapshot_test.go b/snapshot_test.go
index 4192f36..a68016a 100644
--- a/snapshot_test.go
+++ b/snapshot_test.go
@@ -699,17 +699,17 @@ func writeSnapshotFixture(t *testing.T, snapshotDir, environment, fixture string
 }
 
 type snapshotRemoteHandlers struct {
-	authStatus          int
-	authBody            map[string]any
-	snapshotCheckStatus int
-	snapshotCheckBody   map[string]any
+	authStatus           int
+	authBody             map[string]any
+	snapshotCheckStatus  int
+	snapshotCheckBody    map[string]any
 	snapshotCheckRawBody *string
-	resolveStatus       int
-	resolveDomain       map[string]any
-	resolveRawBody      *string
-	resolveUnavailable  bool
-	snapshotChecks      []snapshotCheckStep
-	resolveSteps        []resolveSnapshotStep
+	resolveStatus        int
+	resolveDomain        map[string]any
+	resolveRawBody       *string
+	resolveUnavailable   bool
+	snapshotChecks       []snapshotCheckStep
+	resolveSteps         []resolveSnapshotStep
 }
 
 type snapshotCheckStep struct {
@@ -731,74 +731,90 @@ func newSnapshotTestServer(t *testing.T, handlers snapshotRemoteHandlers) *httpt
 	var mu sync.Mutex
 
 	mux.HandleFunc("/criteria/auth", func(writer http.ResponseWriter, request *http.Request) {
-		assert.Equal(t, http.MethodPost, request.Method)
-		writeJSONResponse(t, writer, handlers.authStatus, handlers.authBody)
+		handleAuth(t, writer, request, handlers)
 	})
 	mux.HandleFunc("/criteria/snapshot_check/", func(writer http.ResponseWriter, request *http.Request) {
-		assert.Equal(t, http.MethodGet, request.Method)
+		handleSnapshotCheck(t, writer, request, handlers, &snapshotCheckIndex, &mu)
+	})
+	mux.HandleFunc("/graphql", func(writer http.ResponseWriter, request *http.Request) {
+		handleGraphQL(t, writer, request, handlers, &resolveIndex, &mu)
+	})
 
-		mu.Lock()
-		defer mu.Unlock()
+	return httptest.NewServer(mux)
+}
 
-		if len(handlers.snapshotChecks) > 0 {
-			step := handlers.snapshotChecks[min(snapshotCheckIndex, len(handlers.snapshotChecks)-1)]
-			snapshotCheckIndex++
-			writeJSONResponse(t, writer, step.status, step.body)
-			return
-		}
+func handleAuth(t *testing.T, writer http.ResponseWriter, request *http.Request, handlers snapshotRemoteHandlers) {
+	assert.Equal(t, http.MethodPost, request.Method)
+	writeJSONResponse(t, writer, handlers.authStatus, handlers.authBody)
+}
 
-		if handlers.snapshotCheckRawBody != nil {
-			writer.Header().Set("Content-Type", "application/json")
-			writer.WriteHeader(handlers.snapshotCheckStatus)
-			_, err := writer.Write([]byte(*handlers.snapshotCheckRawBody))
-			assert.NoError(t, err)
-			return
-		}
+func handleSnapshotCheck(t *testing.T, writer http.ResponseWriter, request *http.Request, handlers snapshotRemoteHandlers, index *int, mu *sync.Mutex) {
+	assert.Equal(t, http.MethodGet, request.Method)
 
-		writeJSONResponse(t, writer, handlers.snapshotCheckStatus, handlers.snapshotCheckBody)
-	})
-	mux.HandleFunc("/graphql", func(writer http.ResponseWriter, request *http.Request) {
-		assert.Equal(t, http.MethodPost, request.Method)
+	mu.Lock()
+	defer mu.Unlock()
 
-		mu.Lock()
-		defer mu.Unlock()
+	if len(handlers.snapshotChecks) > 0 {
+		step := handlers.snapshotChecks[min(*index, len(handlers.snapshotChecks)-1)]
+		*index++
+		writeJSONResponse(t, writer, step.status, step.body)
+		return
+	}
 
-		if handlers.resolveUnavailable {
-			hijacker, ok := writer.(http.Hijacker)
-			assert.True(t, ok)
-			if !ok {
-				return
-			}
+	if handlers.snapshotCheckRawBody != nil {
+		writer.Header().Set("Content-Type", "application/json")
+		writer.WriteHeader(handlers.snapshotCheckStatus)
+		_, err := writer.Write([]byte(*handlers.snapshotCheckRawBody))
+		assert.NoError(t, err)
+		return
+	}
 
-			conn, _, err := hijacker.Hijack()
-			assert.NoError(t, err)
-			if err != nil {
-				return
-			}
+	writeJSONResponse(t, writer, handlers.snapshotCheckStatus, handlers.snapshotCheckBody)
+}
 
-			_ = conn.Close()
-			return
-		}
+func handleGraphQL(t *testing.T, writer http.ResponseWriter, request *http.Request, handlers snapshotRemoteHandlers, index *int, mu *sync.Mutex) {
+	assert.Equal(t, http.MethodPost, request.Method)
 
-		if handlers.resolveRawBody != nil {
-			writer.Header().Set("Content-Type", "application/json")
-			writer.WriteHeader(handlers.resolveStatus)
-			_, err := writer.Write([]byte(*handlers.resolveRawBody))
-			assert.NoError(t, err)
-			return
-		}
+	mu.Lock()
+	defer mu.Unlock()
 
-		if len(handlers.resolveSteps) > 0 {
-			step := handlers.resolveSteps[min(resolveIndex, len(handlers.resolveSteps)-1)]
-			resolveIndex++
-			writeJSONResponse(t, writer, step.status, map[string]any{"data": map[string]any{"domain": step.domain}})
-			return
-		}
+	if handlers.resolveUnavailable {
+		handleUnavailable(t, writer)
+		return
+	}
 
-		writeJSONResponse(t, writer, handlers.resolveStatus, map[string]any{"data": map[string]any{"domain": handlers.resolveDomain}})
-	})
+	if handlers.resolveRawBody != nil {
+		writer.Header().Set("Content-Type", "application/json")
+		writer.WriteHeader(handlers.resolveStatus)
+		_, err := writer.Write([]byte(*handlers.resolveRawBody))
+		assert.NoError(t, err)
+		return
+	}
 
-	return httptest.NewServer(mux)
+	if len(handlers.resolveSteps) > 0 {
+		step := handlers.resolveSteps[min(*index, len(handlers.resolveSteps)-1)]
+		*index++
+		writeJSONResponse(t, writer, step.status, map[string]any{"data": map[string]any{"domain": step.domain}})
+		return
+	}
+
+	writeJSONResponse(t, writer, handlers.resolveStatus, map[string]any{"data": map[string]any{"domain": handlers.resolveDomain}})
+}
+
+func handleUnavailable(t *testing.T, writer http.ResponseWriter) {
+	hijacker, ok := writer.(http.Hijacker)
+	assert.True(t, ok)
+	if !ok {
+		return
+	}
+
+	conn, _, err := hijacker.Hijack()
+	assert.NoError(t, err)
+	if err != nil {
+		return
+	}
+
+	_ = conn.Close()
 }
 
 func min(a, b int) int {