From 5b7f050eb2dc62e7408c715a89bd493c511037cf Mon Sep 17 00:00:00 2001
From: Harish Senthilkumar <harishxr@amazon.com>
Date: Mon, 16 Sep 2024 23:13:22 +0000
Subject: [PATCH] Add integ tests and refactor code

---
 agent/handlers/task_server_setup.go           |  2 +-
 .../handlers/task_server_setup_integ_test.go  | 47 +++++++++++++++++++
 2 files changed, 48 insertions(+), 1 deletion(-)
 create mode 100644 agent/handlers/task_server_setup_integ_test.go

diff --git a/agent/handlers/task_server_setup.go b/agent/handlers/task_server_setup.go
index debfb291bef..89cfba8c6ee 100644
--- a/agent/handlers/task_server_setup.go
+++ b/agent/handlers/task_server_setup.go
@@ -95,7 +95,7 @@ func taskServerSetup(
 		taskProtectionClientFactory, metricsFactory)
 
 	// TODO: Future PR to pass in TMDS server router once all of the handlers have been implemented.
-	registerFaultHandlers(muxRouter, tmdsAgentState, metricsFactory)
+	registerFaultHandlers(nil, tmdsAgentState, metricsFactory)
 
 	return tmds.NewServer(auditLogger,
 		tmds.WithHandler(muxRouter),
diff --git a/agent/handlers/task_server_setup_integ_test.go b/agent/handlers/task_server_setup_integ_test.go
new file mode 100644
index 00000000000..bf5fced9587
--- /dev/null
+++ b/agent/handlers/task_server_setup_integ_test.go
@@ -0,0 +1,47 @@
+//go:build integration
+// +build integration
+
+// Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License"). You may
+// not use this file except in compliance with the License. A copy of the
+// License is located at
+//
+//	http://aws.amazon.com/apache2.0/
+//
+// or in the "license" file accompanying this file. This file is distributed
+// on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+// express or implied. See the License for the specific language governing
+// permissions and limitations under the License.
+
+package handlers
+
+import (
+	"net/http"
+	"testing"
+
+	"github.com/didip/tollbooth"
+	"github.com/gorilla/mux"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestCreateRateLimiter(t *testing.T) {
+	// Start the server
+	go func() {
+		router := mux.NewRouter()
+		router.Handle("/test", tollbooth.LimitFuncHandler(createRateLimiter(), func(w http.ResponseWriter, r *http.Request) {
+			w.WriteHeader(http.StatusOK)
+			w.Write([]byte("Success"))
+		}))
+		http.ListenAndServe(":5932", router)
+	}()
+	// First request should succeed
+	resp, err := http.Get("http://localhost:5932/test")
+	require.NoError(t, err)
+	assert.Equal(t, http.StatusOK, resp.StatusCode)
+	// Second request should be rate-limited (429)
+	resp, err = http.Get("http://localhost:5932/test")
+	require.NoError(t, err)
+	assert.Equal(t, http.StatusTooManyRequests, resp.StatusCode)
+}