Add integ tests and refactor code

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),

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)
+}