From a79fe0a1a0cbb9eb85282325de35dbe33f8a6b1f Mon Sep 17 00:00:00 2001
From: Jason Wilder <jawilder@microsoft.com>
Date: Thu, 23 Jan 2025 20:47:44 -0700
Subject: [PATCH] Increase test code coverage

---
 pkg/http/randomizer_test.go | 38 +++++++++++++++++++++++++++++++++++++
 1 file changed, 38 insertions(+)
 create mode 100644 pkg/http/randomizer_test.go

diff --git a/pkg/http/randomizer_test.go b/pkg/http/randomizer_test.go
new file mode 100644
index 000000000..21ade0993
--- /dev/null
+++ b/pkg/http/randomizer_test.go
@@ -0,0 +1,38 @@
+package http
+
+import (
+	"net/http/httptest"
+	"sync/atomic"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestMaybeCloseConnection(t *testing.T) {
+	// Reset the counter before each test
+	atomic.StoreUint64(&counter, 0)
+
+	t.Run("does not close connection before threshold", func(t *testing.T) {
+		for i := 0; i < closeEvery-1; i++ {
+			req := httptest.NewRequest("GET", "/", nil)
+			resp := httptest.NewRecorder()
+
+			MaybeCloseConnection(resp, req)
+
+			require.False(t, req.Close)
+			require.NotEqual(t, "close", resp.Header().Get("Connection"))
+		}
+	})
+
+	t.Run("closes connection at threshold", func(t *testing.T) {
+		req := httptest.NewRequest("GET", "/", nil)
+		resp := httptest.NewRecorder()
+
+		// Simulate reaching the threshold
+		atomic.StoreUint64(&counter, closeEvery-1)
+		MaybeCloseConnection(resp, req)
+
+		require.True(t, req.Close)
+		require.Equal(t, "close", resp.Header().Get("Connection"))
+	})
+}