From 7405f66036537095b52c277d9b56969df33bfa57 Mon Sep 17 00:00:00 2001
From: George Goldberg <george@gberg.me>
Date: Tue, 5 Sep 2017 20:34:17 +0100
Subject: PLT-7519: Better rate-limiting. (#7365)

---
 utils/utils_test.go | 55 +++++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 55 insertions(+)

(limited to 'utils/utils_test.go')

diff --git a/utils/utils_test.go b/utils/utils_test.go
index b80247867..b18566786 100644
--- a/utils/utils_test.go
+++ b/utils/utils_test.go
@@ -4,7 +4,10 @@
 package utils
 
 import (
+	"net/http"
 	"testing"
+
+	"github.com/stretchr/testify/assert"
 )
 
 func TestStringArrayIntersection(t *testing.T) {
@@ -44,3 +47,55 @@ func TestRemoveDuplicatesFromStringArray(t *testing.T) {
 		t.Fatal("should be 3")
 	}
 }
+
+func TestGetIpAddress(t *testing.T) {
+	// Test with a single IP in the X-Forwarded-For
+	httpRequest1 := http.Request{
+		Header: http.Header{
+			"X-Forwarded-For": []string{"10.0.0.1"},
+			"X-Real-Ip":       []string{"10.1.0.1"},
+		},
+		RemoteAddr: "10.2.0.1:12345",
+	}
+
+	assert.Equal(t, "10.0.0.1", GetIpAddress(&httpRequest1))
+
+	// Test with multiple IPs in the X-Forwarded-For
+	httpRequest2 := http.Request{
+		Header: http.Header{
+			"X-Forwarded-For": []string{"10.0.0.1,  10.0.0.2, 10.0.0.3"},
+			"X-Real-Ip":       []string{"10.1.0.1"},
+		},
+		RemoteAddr: "10.2.0.1:12345",
+	}
+
+	assert.Equal(t, "10.0.0.1", GetIpAddress(&httpRequest2))
+
+	// Test with an empty X-Forwarded-For
+	httpRequest3 := http.Request{
+		Header: http.Header{
+			"X-Forwarded-For": []string{""},
+			"X-Real-Ip":       []string{"10.1.0.1"},
+		},
+		RemoteAddr: "10.2.0.1:12345",
+	}
+
+	assert.Equal(t, "10.1.0.1", GetIpAddress(&httpRequest3))
+
+	// Test without an X-Fowarded-For
+	httpRequest4 := http.Request{
+		Header: http.Header{
+			"X-Real-Ip": []string{"10.1.0.1"},
+		},
+		RemoteAddr: "10.2.0.1:12345",
+	}
+
+	assert.Equal(t, "10.1.0.1", GetIpAddress(&httpRequest4))
+
+	// Test without any headers
+	httpRequest5 := http.Request{
+		RemoteAddr: "10.2.0.1:12345",
+	}
+
+	assert.Equal(t, "10.2.0.1", GetIpAddress(&httpRequest5))
+}
-- 
cgit v1.2.3-1-g7c22