fix broken tests in ShouldBlockRequest

Author: bitterpanda63

agent_api/src/test/java/ShouldBlockRequestTest.java

@@ -4,6 +4,7 @@
 import dev.aikido.agent_api.context.Context;
 import dev.aikido.agent_api.context.ContextObject;
 import dev.aikido.agent_api.context.User;
+import dev.aikido.agent_api.storage.RateLimiterStore;
 import dev.aikido.agent_api.storage.ServiceConfigStore;
 import dev.aikido.agent_api.storage.statistics.StatisticsStore;
 import org.junit.jupiter.api.AfterEach;
@@ -40,12 +41,15 @@ public static void clean() {
         Context.set(null);
         ServiceConfigStore.updateFromAPIResponse(emptyAPIResponse);
         StatisticsStore.clear();
+        RateLimiterStore.clear();
     };
 
     @AfterEach
     public void tearDown() throws SQLException {
         Context.set(null);
         ServiceConfigStore.updateFromAPIResponse(emptyAPIResponse);
+        StatisticsStore.clear();
+        RateLimiterStore.clear();
     }
 
     @Test
@@ -135,22 +139,39 @@ public void testEndpointsExistButNoMatch() throws SQLException {
 
     @Test
     public void testEndpointsExistWithMatch() throws SQLException {
-        Context.set(new SampleContextObject());
+        Context.set(null);
         setEmptyConfigWithEndpointList(List.of(
-                new Endpoint("GET", "/api/*", 1, 1000, Collections.emptyList(), false, false, false)
+            new Endpoint("GET", "/api/*", 1, 1000, Collections.emptyList(), false, false, false)
         ));
 
         // Test with match & rate-limiting disabled :
         var res1 = ShouldBlockRequest.shouldBlockRequest();
         assertFalse(res1.block());
 
+        Context.set(null);
         setEmptyConfigWithEndpointList(List.of(
-                new Endpoint("GET", "/api/*", 1, 1000, Collections.emptyList(), false, false, true)
+            new Endpoint("GET", "/api/*", 1, 1000, Collections.emptyList(), false, false, true)
         ));
 
         // Test with match & rate-limiting enabled :
         var res2 = ShouldBlockRequest.shouldBlockRequest();
         assertFalse(res2.block());
+    }
+
+    @Test
+    public void testEndpointsExistAndGetsRateLimited() throws SQLException {
+        setEmptyConfigWithEndpointList(List.of(
+                new Endpoint("GET", "/api/*", 2, 1000, Collections.emptyList(), false, false, true)
+        ));
+
+        // Test with match
+        Context.set(new SampleContextObject());
+        var res1 = ShouldBlockRequest.shouldBlockRequest();
+        assertFalse(res1.block());
+
+        // Test with match
+        var res2 = ShouldBlockRequest.shouldBlockRequest();
+        assertFalse(res2.block());
         assertEquals(0, StatisticsStore.getStatsRecord().requests().rateLimited());
 
         var res3 = ShouldBlockRequest.shouldBlockRequest();