Adding quota tests

pom.xml

@@ -166,10 +166,8 @@
                     <destFile>${basedir}/target/coverage-reports/jacoco-unit.exec</destFile>
                     <dataFile>${basedir}/target/coverage-reports/jacoco-unit.exec</dataFile>
                     <excludes>
-                        <exclude>**/properties/**</exclude>
-                        <exclude>**/*Configuration.*</exclude>
-                        <exclude>**/*Application.*</exclude>
-                        <exclude>**/*Rate.*</exclude>
+                        <exclude>**/test/**</exclude>
+                        <exclude>**/tests/**</exclude>
                     </excludes>
                 </configuration>
                 <executions>

spring-cloud-zuul-ratelimit-core/src/main/java/com/marcosbarbero/cloud/autoconfigure/zuul/ratelimit/filters/RateLimitPostFilter.java

@@ -16,16 +16,13 @@
 
 package com.marcosbarbero.cloud.autoconfigure.zuul.ratelimit.filters;
 
-import static java.util.concurrent.TimeUnit.MILLISECONDS;
 import static org.springframework.web.context.request.RequestAttributes.SCOPE_REQUEST;
 
-import com.marcosbarbero.cloud.autoconfigure.zuul.ratelimit.config.Rate;
 import com.marcosbarbero.cloud.autoconfigure.zuul.ratelimit.config.RateLimitKeyGenerator;
 import com.marcosbarbero.cloud.autoconfigure.zuul.ratelimit.config.RateLimiter;
 import com.marcosbarbero.cloud.autoconfigure.zuul.ratelimit.config.properties.RateLimitProperties;
 import com.netflix.zuul.context.RequestContext;
 import javax.servlet.http.HttpServletRequest;
-import javax.servlet.http.HttpServletResponse;
 import org.springframework.cloud.netflix.zuul.filters.Route;
 import org.springframework.cloud.netflix.zuul.filters.RouteLocator;
 import org.springframework.cloud.netflix.zuul.filters.support.FilterConstants;
@@ -75,23 +72,14 @@ private Long getRequestStartTime() {
 
     public Object run() {
         final RequestContext ctx = RequestContext.getCurrentContext();
-        final HttpServletResponse response = ctx.getResponse();
         final HttpServletRequest request = ctx.getRequest();
         final Route route = route();
 
         policy(route).ifPresent(policy -> {
 
             final Long requestTime = System.currentTimeMillis() - getRequestStartTime();
             final String key = rateLimitKeyGenerator.key(request, route, policy);
-            final Rate rate = rateLimiter.consume(policy, key, requestTime);
-
-            final Long quota = policy.getQuota();
-            final Long remainingQuota = rate.getRemainingQuota();
-            if (quota != null) {
-                RequestContextHolder.getRequestAttributes().setAttribute(REQUEST_START_TIME, System.currentTimeMillis(), SCOPE_REQUEST);
-                response.setHeader(QUOTA_HEADER, String.valueOf(quota));
-                response.setHeader(REMAINING_QUOTA_HEADER, String.valueOf(MILLISECONDS.toSeconds(Math.max(remainingQuota, 0))));
-            }
+            rateLimiter.consume(policy, key, requestTime);
         });
 
         return null;

spring-cloud-zuul-ratelimit-core/src/main/java/com/marcosbarbero/cloud/autoconfigure/zuul/ratelimit/filters/RateLimitPreFilter.java

@@ -16,6 +16,7 @@
 
 package com.marcosbarbero.cloud.autoconfigure.zuul.ratelimit.filters;
 
+import static java.util.concurrent.TimeUnit.MILLISECONDS;
 import static org.springframework.web.context.request.RequestAttributes.SCOPE_REQUEST;
 
 import com.marcosbarbero.cloud.autoconfigure.zuul.ratelimit.config.Rate;
@@ -85,6 +86,8 @@ public Object run() {
             final Long remainingQuota = rate.getRemainingQuota();
             if (quota != null) {
                 RequestContextHolder.getRequestAttributes().setAttribute(REQUEST_START_TIME, System.currentTimeMillis(), SCOPE_REQUEST);
+                response.setHeader(QUOTA_HEADER, String.valueOf(quota));
+                response.setHeader(REMAINING_QUOTA_HEADER, String.valueOf(MILLISECONDS.toSeconds(Math.max(remainingQuota, 0))));
             }
 
             response.setHeader(RESET_HEADER, String.valueOf(rate.getReset()));

spring-cloud-zuul-ratelimit-core/src/test/java/com/marcosbarbero/cloud/autoconfigure/zuul/ratelimit/config/repository/ConsulRateLimiterTest.java

@@ -1,23 +1,34 @@
 package com.marcosbarbero.cloud.autoconfigure.zuul.ratelimit.config.repository;
 
+import static org.assertj.core.api.Assertions.assertThat;
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.verifyZeroInteractions;
 import static org.mockito.Mockito.when;
 
 import com.ecwid.consul.v1.ConsulClient;
 import com.ecwid.consul.v1.Response;
 import com.ecwid.consul.v1.kv.model.GetValue;
+import com.fasterxml.jackson.core.JsonProcessingException;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.google.common.collect.Maps;
+import com.marcosbarbero.cloud.autoconfigure.zuul.ratelimit.config.Rate;
+import java.io.IOException;
 import java.util.Base64;
 import java.util.Map;
 import org.junit.Before;
+import org.junit.Test;
 import org.mockito.Mock;
+import org.mockito.Mockito;
 import org.mockito.MockitoAnnotations;
 
 public class ConsulRateLimiterTest extends BaseRateLimiterTest {
 
     @Mock
     private ConsulClient consulClient;
+    @Mock
+    private ObjectMapper objectMapper;
 
     @Before
     public void setUp() {
@@ -39,4 +50,23 @@ public void setUp() {
         ObjectMapper objectMapper = new ObjectMapper();
         target = new ConsulRateLimiter(consulClient, objectMapper);
     }
+
+    @Test
+    public void testGetRateException() throws IOException {
+        when(objectMapper.readValue(anyString(), eq(Rate.class))).thenThrow(new IOException());
+        ConsulRateLimiter consulRateLimiter = new ConsulRateLimiter(consulClient, objectMapper);
+
+        Rate rate = consulRateLimiter.getRate("");
+        assertThat(rate).isNull();
+    }
+
+    @Test
+    public void testSaveRateException() throws IOException {
+        JsonProcessingException jsonProcessingException = Mockito.mock(JsonProcessingException.class);
+        when(objectMapper.writeValueAsString(any(Rate.class))).thenThrow(jsonProcessingException);
+        ConsulRateLimiter consulRateLimiter = new ConsulRateLimiter(consulClient, objectMapper);
+
+        consulRateLimiter.saveRate(null);
+        verifyZeroInteractions(consulClient);
+    }
 }

spring-cloud-zuul-ratelimit-core/src/test/java/com/marcosbarbero/cloud/autoconfigure/zuul/ratelimit/filters/post/RateLimitPostFilterTest.java

@@ -1,6 +1,11 @@
 package com.marcosbarbero.cloud.autoconfigure.zuul.ratelimit.filters.post;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyLong;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyZeroInteractions;
 import static org.mockito.Mockito.when;
 import static org.springframework.web.context.request.RequestAttributes.SCOPE_REQUEST;
 
@@ -87,10 +92,29 @@ public void testShouldFilterOnNullStartTime() {
     @Test
     public void testShouldFilter() {
         rateLimitProperties.setEnabled(true);
-        when(requestAttributes.getAttribute(RateLimitPreFilter.REQUEST_START_TIME, SCOPE_REQUEST)).thenReturn(5L);
+        when(requestAttributes.getAttribute(RateLimitPreFilter.REQUEST_START_TIME, SCOPE_REQUEST)).thenReturn(System.currentTimeMillis());
         Policy defaultPolicy = new Policy();
         rateLimitProperties.setDefaultPolicy(defaultPolicy);
 
         assertThat(target.shouldFilter()).isEqualTo(true);
     }
+
+    @Test
+    public void testRunNoPolicy() {
+        target.run();
+        verifyZeroInteractions(rateLimiter);
+    }
+
+    @Test
+    public void testRun() {
+        rateLimitProperties.setEnabled(true);
+        when(requestAttributes.getAttribute(RateLimitPreFilter.REQUEST_START_TIME, SCOPE_REQUEST)).thenReturn(System.currentTimeMillis());
+        Policy defaultPolicy = new Policy();
+        defaultPolicy.setQuota(2L);
+        rateLimitProperties.setDefaultPolicy(defaultPolicy);
+        when(rateLimitKeyGenerator.key(any(), any(), any())).thenReturn("generatedKey");
+
+        target.run();
+        verify(rateLimiter).consume(eq(defaultPolicy), eq("generatedKey"), anyLong());
+    }
 }

spring-cloud-zuul-ratelimit-core/src/test/java/com/marcosbarbero/cloud/autoconfigure/zuul/ratelimit/filters/pre/BaseRateLimitPreFilterTest.java

@@ -21,11 +21,15 @@
 import java.util.concurrent.TimeUnit;
 import org.junit.Before;
 import org.junit.Test;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
 import org.springframework.cloud.netflix.zuul.filters.Route;
 import org.springframework.cloud.netflix.zuul.filters.RouteLocator;
 import org.springframework.cloud.netflix.zuul.metrics.EmptyCounterFactory;
 import org.springframework.mock.web.MockHttpServletRequest;
 import org.springframework.mock.web.MockHttpServletResponse;
+import org.springframework.web.context.request.RequestAttributes;
+import org.springframework.web.context.request.RequestContextHolder;
 import org.springframework.web.util.UrlPathHelper;
 
 /**
@@ -36,6 +40,9 @@ public abstract class BaseRateLimitPreFilterTest {
 
     RateLimitPreFilter filter;
 
+    @Mock
+    RequestAttributes requestAttributes;
+
     MockHttpServletRequest request;
     MockHttpServletResponse response;
 
@@ -56,6 +63,7 @@ private RateLimitProperties properties() {
 
         Policy policy = new Policy();
         policy.setLimit(2L);
+        policy.setQuota(2L);
         policy.setRefreshInterval(2L);
         policy.setType(asList(Policy.Type.ORIGIN, Policy.Type.URL, Policy.Type.USER));
 
@@ -76,6 +84,7 @@ void setRateLimiter(RateLimiter rateLimiter) {
 
     @Before
     public void setUp() {
+        MockitoAnnotations.initMocks(this);
         CounterFactory.initialize(new EmptyCounterFactory());
         this.request = new MockHttpServletRequest();
         this.response = new MockHttpServletResponse();
@@ -84,6 +93,7 @@ public void setUp() {
         this.filter = new RateLimitPreFilter(this.properties(), this.routeLocator(), urlPathHelper, this.rateLimiter, this.rateLimitKeyGenerator);
         this.context = new RequestContext();
         RequestContext.testSetCurrentContext(this.context);
+        RequestContextHolder.setRequestAttributes(requestAttributes);
         this.context.clear();
         this.context.setRequest(this.request);
         this.context.setResponse(this.response);

spring-cloud-zuul-ratelimit-core/src/test/java/com/marcosbarbero/cloud/autoconfigure/zuul/ratelimit/filters/pre/ConsulRateLimitPreFilterTest.java

@@ -29,7 +29,7 @@ public class ConsulRateLimitPreFilterTest extends BaseRateLimitPreFilterTest {
     private ObjectMapper objectMapper = new ObjectMapper();
 
     private Rate rate(long remaining) {
-        return new Rate("key", remaining, null, 100L, new Date(System.currentTimeMillis() + SECONDS.toMillis(2)));
+        return new Rate("key", remaining, 2000L, 100L, new Date(System.currentTimeMillis() + SECONDS.toMillis(2)));
     }
 
     @Before

spring-cloud-zuul-ratelimit-tests/consul/src/main/java/com/marcosbarbero/tests/ConsulApplication.java

@@ -12,7 +12,6 @@
 import org.springframework.http.ResponseEntity;
 import org.springframework.web.bind.annotation.GetMapping;
 import org.springframework.web.bind.annotation.PathVariable;
-import org.springframework.web.bind.annotation.RequestMapping;
 import org.springframework.web.bind.annotation.RestController;
 
 /**
@@ -28,7 +27,6 @@ public static void main(String... args) {
     }
 
     @RestController
-    @RequestMapping("/services")
     public class ServiceController {
 
         public static final String RESPONSE_BODY = "ResponseBody";
@@ -52,6 +50,12 @@ public ResponseEntity<String> serviceC() {
         public ResponseEntity<String> serviceD(@PathVariable String paramName) {
             return ResponseEntity.ok(RESPONSE_BODY + " " + paramName);
         }
+
+        @GetMapping("/serviceE")
+        public ResponseEntity<String> serviceE() throws InterruptedException {
+            Thread.sleep(1100);
+            return ResponseEntity.ok(RESPONSE_BODY);
+        }
     }
 
 

spring-cloud-zuul-ratelimit-tests/consul/src/main/resources/application.yml

@@ -2,17 +2,20 @@ zuul:
   routes:
     serviceA:
       path: /serviceA
-      url: forward:/services
+      url: forward:/
     serviceB:
       path: /serviceB
-      url: forward:/services
+      url: forward:/
     serviceC:
       path: /serviceC
-      url: forward:/services
+      url: forward:/
     serviceD:
       strip-prefix: false
       path: /serviceD/**
-      url: forward:/services
+      url: forward:/
+    serviceE:
+      path: /serviceE
+      url: forward:/
   ratelimit:
     enabled: true
     repository: CONSUL
@@ -32,4 +35,9 @@ zuul:
         refresh-interval: 60
         type:
           - url
+      serviceE:
+        quota: 1
+        refresh-interval: 60
+        type:
+          - origin
   strip-prefix: true

spring-cloud-zuul-ratelimit-tests/consul/src/test/java/com/marcosbarbero/tests/it/ConsulApplicationTestIT.java

@@ -48,15 +48,15 @@ public void testConsulRateLimiter() {
     public void testNotExceedingCapacityRequest() {
         ResponseEntity<String> response = this.restTemplate.getForEntity("/serviceA", String.class);
         HttpHeaders headers = response.getHeaders();
-        assertHeaders(headers, false);
+        assertHeaders(headers, false, false);
         assertEquals(OK, response.getStatusCode());
     }
 
     @Test
     public void testExceedingCapacity() throws InterruptedException {
         ResponseEntity<String> response = this.restTemplate.getForEntity("/serviceB", String.class);
         HttpHeaders headers = response.getHeaders();
-        assertHeaders(headers, false);
+        assertHeaders(headers, false, false);
         assertEquals(OK, response.getStatusCode());
 
         for (int i = 0; i < 2; i++) {
@@ -70,15 +70,15 @@ public void testExceedingCapacity() throws InterruptedException {
 
         response = this.restTemplate.getForEntity("/serviceB", String.class);
         headers = response.getHeaders();
-        assertHeaders(headers, false);
+        assertHeaders(headers, false, false);
         assertEquals(OK, response.getStatusCode());
     }
 
     @Test
     public void testNoRateLimit() {
         ResponseEntity<String> response = this.restTemplate.getForEntity("/serviceC", String.class);
         HttpHeaders headers = response.getHeaders();
-        assertHeaders(headers, true);
+        assertHeaders(headers, true, false);
         assertEquals(OK, response.getStatusCode());
     }
 
@@ -93,25 +93,49 @@ public void testMultipleUrls() {
 
             ResponseEntity<String> response = this.restTemplate.getForEntity("/serviceD/" + randomPath, String.class);
             HttpHeaders headers = response.getHeaders();
-            assertHeaders(headers, false);
+            assertHeaders(headers, false, false);
             assertEquals(OK, response.getStatusCode());
         }
     }
 
-    private void assertHeaders(HttpHeaders headers, boolean nullable) {
+    @Test
+    public void testExceedingQuotaCapacityRequest() {
+        ResponseEntity<String> response = this.restTemplate.getForEntity("/serviceE", String.class);
+        HttpHeaders headers = response.getHeaders();
+        assertHeaders(headers, false, true);
+        assertEquals(OK, response.getStatusCode());
+
+        response = this.restTemplate.getForEntity("/serviceE", String.class);
+        headers = response.getHeaders();
+        assertHeaders(headers, false, true);
+        assertEquals(TOO_MANY_REQUESTS, response.getStatusCode());
+    }
+
+    private void assertHeaders(HttpHeaders headers, boolean nullable, boolean quotaHeaders) {
+        String quota = headers.getFirst(RateLimitPreFilter.QUOTA_HEADER);
+        String remainingQuota = headers.getFirst(RateLimitPreFilter.REMAINING_QUOTA_HEADER);
         String limit = headers.getFirst(RateLimitPreFilter.LIMIT_HEADER);
         String remaining = headers.getFirst(RateLimitPreFilter.REMAINING_HEADER);
         String reset = headers.getFirst(RateLimitPreFilter.RESET_HEADER);
 
         if (nullable) {
-            assertNull(limit);
-            assertNull(remaining);
+            if (quotaHeaders) {
+                assertNull(quota);
+                assertNull(remainingQuota);
+            } else {
+                assertNull(limit);
+                assertNull(remaining);
+            }
             assertNull(reset);
         } else {
-            assertNotNull(limit);
-            assertNotNull(remaining);
+            if (quotaHeaders) {
+                assertNotNull(quota);
+                assertNotNull(remainingQuota);
+            } else {
+                assertNotNull(limit);
+                assertNotNull(remaining);
+            }
             assertNotNull(reset);
         }
     }
-
 }