From ce9e9f3e0d2ee8be9d946dfc269f5fb524fb6d7d Mon Sep 17 00:00:00 2001
From: Philipp Heckel <pheckel@datto.com>
Date: Mon, 14 Feb 2022 19:23:57 -0500
Subject: [PATCH] Add test for IP exempt setting

---
 server/server_test.go | 37 +++++++++++++++++++++++++++++++++++++
 1 file changed, 37 insertions(+)

diff --git a/server/server_test.go b/server/server_test.go
index 2e40bd44..614cd5c9 100644
--- a/server/server_test.go
+++ b/server/server_test.go
@@ -690,6 +690,43 @@ func (t *testMailer) Send(from, to string, m *message) error {
 	return nil
 }
 
+func TestServer_PublishTooRequests_Defaults(t *testing.T) {
+	s := newTestServer(t, newTestConfig(t))
+	for i := 0; i < 60; i++ {
+		response := request(t, s, "PUT", "/mytopic", fmt.Sprintf("message %d", i), nil)
+		require.Equal(t, 200, response.Code)
+	}
+	response := request(t, s, "PUT", "/mytopic", "message", nil)
+	require.Equal(t, 429, response.Code)
+}
+
+func TestServer_PublishTooRequests_Defaults_ExemptHosts(t *testing.T) {
+	c := newTestConfig(t)
+	c.VisitorRequestExemptIPAddrs = []string{"9.9.9.9"} // see request()
+	s := newTestServer(t, c)
+	for i := 0; i < 65; i++ { // > 60
+		response := request(t, s, "PUT", "/mytopic", fmt.Sprintf("message %d", i), nil)
+		require.Equal(t, 200, response.Code)
+	}
+}
+
+func TestServer_PublishTooRequests_ShortReplenish(t *testing.T) {
+	c := newTestConfig(t)
+	c.VisitorRequestLimitBurst = 60
+	c.VisitorRequestLimitReplenish = 500 * time.Millisecond
+	s := newTestServer(t, c)
+	for i := 0; i < 60; i++ {
+		response := request(t, s, "PUT", "/mytopic", fmt.Sprintf("message %d", i), nil)
+		require.Equal(t, 200, response.Code)
+	}
+	response := request(t, s, "PUT", "/mytopic", "message", nil)
+	require.Equal(t, 429, response.Code)
+
+	time.Sleep(510 * time.Millisecond)
+	response = request(t, s, "PUT", "/mytopic", "message", nil)
+	require.Equal(t, 200, response.Code)
+}
+
 func TestServer_PublishTooManyEmails_Defaults(t *testing.T) {
 	s := newTestServer(t, newTestConfig(t))
 	s.mailer = &testMailer{}