1 files changed, 49 insertions, 0 deletions
diff --git a/workhorse/internal/channel/channel_test.go b/workhorse/internal/channel/channel_test.go
new file mode 100644
index 00000000000..fade6e42c27
--- /dev/null
+++ b/workhorse/internal/channel/channel_test.go
@@ -0,0 +1,49 @@
+package channel
+
+import (
+	"net/http"
+	"testing"
+)
+
+func TestSetForwardedForGeneratesHeader(t *testing.T) {
+	testCases := []struct {
+		remoteAddr           string
+		previousForwardedFor []string
+		expected             string
+	}{
+		{
+			"8.8.8.8:3000",
+			nil,
+			"8.8.8.8",
+		},
+		{
+			"8.8.8.8:3000",
+			[]string{"138.124.33.63, 151.146.211.237"},
+			"138.124.33.63, 151.146.211.237, 8.8.8.8",
+		},
+		{
+			"8.8.8.8:3000",
+			[]string{"8.154.76.107", "115.206.118.179"},
+			"8.154.76.107, 115.206.118.179, 8.8.8.8",
+		},
+	}
+	for _, tc := range testCases {
+		headers := http.Header{}
+		originalRequest := http.Request{
+			RemoteAddr: tc.remoteAddr,
+		}
+
+		if tc.previousForwardedFor != nil {
+			originalRequest.Header = http.Header{
+				"X-Forwarded-For": tc.previousForwardedFor,
+			}
+		}
+
+		setForwardedFor(&headers, &originalRequest)
+
+		result := headers.Get("X-Forwarded-For")
+		if result != tc.expected {
+			t.Fatalf("Expected %v, got %v", tc.expected, result)
+		}
+	}
+}