Update test cases in ShouldBlockRequest

Author: bitterpanda63

agent_api/src/test/java/ShouldBlockRequestTest.java

@@ -4,7 +4,10 @@
 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.routes.RoutesStore;
+import dev.aikido.agent_api.storage.statistics.StatisticsStore;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.Test;
@@ -38,12 +41,18 @@ public SampleContextObject() {
     public static void clean() {
         Context.set(null);
         ServiceConfigStore.updateFromAPIResponse(emptyAPIResponse);
+        StatisticsStore.clear();
+        RoutesStore.clear();
+        RateLimiterStore.clear();
     };
 
     @AfterEach
     public void tearDown() throws SQLException {
         Context.set(null);
         ServiceConfigStore.updateFromAPIResponse(emptyAPIResponse);
+        StatisticsStore.clear();
+        RoutesStore.clear();
+        RateLimiterStore.clear();
     }
 
     @Test
@@ -59,6 +68,7 @@ public void testNoContext() throws SQLException {
         // Test with thread cache not set :
         var res2 = ShouldBlockRequest.shouldBlockRequest();
         assertFalse(res2.block());
+        assertEquals(0, StatisticsStore.getStatsRecord().requests().rateLimited());
     }
 
     @Test
@@ -112,7 +122,8 @@ public void testUserSet() throws SQLException {
 
     @Test
     public void testEndpointsExistButNoMatch() throws SQLException {
-        Context.set(null);
+        ContextObject ctx = new SampleContextObject();
+        Context.set(ctx);
         setEmptyConfigWithEndpointList(List.of(
                 new Endpoint("POST", "/api2/*", 1, 1000, Collections.emptyList(), false, false, false)
         ));
@@ -121,7 +132,6 @@ public void testEndpointsExistButNoMatch() throws SQLException {
         var res1 = ShouldBlockRequest.shouldBlockRequest();
         assertFalse(res1.block());
 
-        Context.set(null);
         setEmptyConfigWithEndpointList(List.of(
                 new Endpoint("POST", "/api2/*", 1, 1000, Collections.emptyList(), false, false, true)
         ));
@@ -133,7 +143,8 @@ public void testEndpointsExistButNoMatch() throws SQLException {
 
     @Test
     public void testEndpointsExistWithMatch() throws SQLException {
-        Context.set(null);
+        ContextObject ctx = new SampleContextObject();
+        Context.set(ctx);
         setEmptyConfigWithEndpointList(List.of(
                 new Endpoint("GET", "/api/*", 1, 1000, Collections.emptyList(), false, false, false)
         ));
@@ -142,14 +153,26 @@ public void testEndpointsExistWithMatch() throws SQLException {
         var res1 = ShouldBlockRequest.shouldBlockRequest();
         assertFalse(res1.block());
 
-        Context.set(null);
         setEmptyConfigWithEndpointList(List.of(
                 new Endpoint("GET", "/api/*", 1, 1000, Collections.emptyList(), false, false, true)
         ));
 
         // Test with match & rate-limiting enabled :
         var res2 = ShouldBlockRequest.shouldBlockRequest();
         assertFalse(res2.block());
+        assertEquals(0, StatisticsStore.getStatsRecord().requests().rateLimited());
+
+
+        var res3 = ShouldBlockRequest.shouldBlockRequest();
+        var res4 = ShouldBlockRequest.shouldBlockRequest();
+        assertTrue(res3.block());
+        assertTrue(res4.block());
+        assertEquals("ip", res3.data().trigger());
+        assertEquals("192.168.1.1", res3.data().ip());
+        assertEquals("ratelimited", res3.data().type());
+        assertEquals(2, StatisticsStore.getStatsRecord().requests().rateLimited());
+        assertEquals(2, RoutesStore.getRoutesAsList()[0].getRateLimitCount());
+
     }
 
     @Test