Merge pull request #9279 from priyank-sriv/bael-3982

BAEL-3982 Rate Limiting Spring API using Bucket4j

spring-boot-modules/spring-boot-libraries/pom.xml

@@ -87,7 +87,35 @@
             <artifactId>javase</artifactId>
             <version>${zxing.version}</version>
         </dependency>
-
+        <!-- Bucket4j -->
+        <dependency>
+            <groupId>com.github.vladimir-bukhtoyarov</groupId>
+            <artifactId>bucket4j-core</artifactId>
+            <version>${bucket4j.version}</version>
+        </dependency>
+        <dependency>
+            <groupId>com.giffing.bucket4j.spring.boot.starter</groupId>
+            <artifactId>bucket4j-spring-boot-starter</artifactId>
+            <version>${bucket4j-spring-boot-starter.version}</version>
+        </dependency>
+        <dependency>
+            <groupId>org.springframework.boot</groupId>
+            <artifactId>spring-boot-starter-cache</artifactId>
+            </dependency>
+        <dependency>
+            <groupId>javax.cache</groupId>
+            <artifactId>cache-api</artifactId>
+            </dependency>
+        <dependency>
+            <groupId>com.github.ben-manes.caffeine</groupId>
+            <artifactId>caffeine</artifactId>
+            <version>${caffeine.version}</version>
+        </dependency>
+        <dependency>
+            <groupId>com.github.ben-manes.caffeine</groupId>
+            <artifactId>jcache</artifactId>
+            <version>${caffeine.version}</version>
+        </dependency>
     </dependencies>
 
     <repositories>
@@ -200,6 +228,9 @@
         <barcode4j.version>2.1</barcode4j.version>
         <qrgen.version>2.6.0</qrgen.version>
         <zxing.version>3.3.0</zxing.version>
+        <bucket4j.version>4.10.0</bucket4j.version>
+        <bucket4j-spring-boot-starter.version>0.2.0</bucket4j-spring-boot-starter.version>
+        <caffeine.version>2.8.2</caffeine.version>
     </properties>
 
 </project>

spring-boot-modules/spring-boot-libraries/src/main/java/com/baeldung/ratelimiting/bootstarterapp/Bucket4jRateLimitApp.java

@@ -0,0 +1,21 @@
+package com.baeldung.ratelimiting.bootstarterapp;
+
+import org.springframework.boot.autoconfigure.SpringBootApplication;
+import org.springframework.boot.autoconfigure.jdbc.DataSourceAutoConfiguration;
+import org.springframework.boot.autoconfigure.security.servlet.SecurityAutoConfiguration;
+import org.springframework.boot.builder.SpringApplicationBuilder;
+import org.springframework.cache.annotation.EnableCaching;
+
+@SpringBootApplication(scanBasePackages = "com.baeldung.ratelimiting", exclude = {
+    DataSourceAutoConfiguration.class,
+    SecurityAutoConfiguration.class,
+})
+@EnableCaching
+public class Bucket4jRateLimitApp {
+
+  public static void main(String[] args) {
+    new SpringApplicationBuilder(Bucket4jRateLimitApp.class)
+        .properties("spring.config.location=classpath:ratelimiting/application-bucket4j-starter.yml")
+        .run(args);
+  }
+}

spring-boot-modules/spring-boot-libraries/src/main/java/com/baeldung/ratelimiting/bucket4japp/Bucket4jRateLimitApp.java

@@ -0,0 +1,35 @@
+package com.baeldung.ratelimiting.bucket4japp;
+
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.boot.autoconfigure.SpringBootApplication;
+import org.springframework.boot.autoconfigure.jdbc.DataSourceAutoConfiguration;
+import org.springframework.boot.autoconfigure.security.servlet.SecurityAutoConfiguration;
+import org.springframework.boot.builder.SpringApplicationBuilder;
+import org.springframework.context.annotation.Lazy;
+import org.springframework.web.servlet.config.annotation.InterceptorRegistry;
+import org.springframework.web.servlet.config.annotation.WebMvcConfigurer;
+
+import com.baeldung.ratelimiting.bucket4japp.interceptor.RateLimitInterceptor;
+
+@SpringBootApplication(scanBasePackages = "com.baeldung.ratelimiting", exclude = { 
+        DataSourceAutoConfiguration.class,
+        SecurityAutoConfiguration.class
+})
+public class Bucket4jRateLimitApp implements WebMvcConfigurer {
+    
+    @Autowired
+    @Lazy
+    private RateLimitInterceptor interceptor;
+
+    @Override
+    public void addInterceptors(InterceptorRegistry registry) {
+        registry.addInterceptor(interceptor)
+            .addPathPatterns("/api/v1/area/**");
+    }
+
+    public static void main(String[] args) {
+        new SpringApplicationBuilder(Bucket4jRateLimitApp.class)
+            .properties("spring.config.location=classpath:ratelimiting/application-bucket4j.yml")
+            .run(args);
+    }
+}

spring-boot-modules/spring-boot-libraries/src/main/java/com/baeldung/ratelimiting/bucket4japp/interceptor/RateLimitInterceptor.java

@@ -0,0 +1,57 @@
+package com.baeldung.ratelimiting.bucket4japp.interceptor;
+
+import javax.servlet.http.HttpServletRequest;
+import javax.servlet.http.HttpServletResponse;
+
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.http.HttpStatus;
+import org.springframework.http.MediaType;
+import org.springframework.stereotype.Component;
+import org.springframework.web.servlet.HandlerInterceptor;
+
+import com.baeldung.ratelimiting.bucket4japp.service.PricingPlanService;
+
+import io.github.bucket4j.Bucket;
+import io.github.bucket4j.ConsumptionProbe;
+
+@Component
+public class RateLimitInterceptor implements HandlerInterceptor {
+
+    private static final String HEADER_API_KEY = "X-api-key";
+    private static final String HEADER_LIMIT_REMAINING = "X-Rate-Limit-Remaining";
+    private static final String HEADER_RETRY_AFTER = "X-Rate-Limit-Retry-After-Seconds";
+
+    @Autowired
+    private PricingPlanService pricingPlanService;
+
+    @Override
+    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) throws Exception {
+
+        String apiKey = request.getHeader(HEADER_API_KEY);
+
+        if (apiKey == null || apiKey.isEmpty()) {
+            response.sendError(HttpStatus.BAD_REQUEST.value(), "Missing Header: " + HEADER_API_KEY);
+            return false;
+        }
+
+        Bucket tokenBucket = pricingPlanService.resolveBucket(apiKey);
+
+        ConsumptionProbe probe = tokenBucket.tryConsumeAndReturnRemaining(1);
+
+        if (probe.isConsumed()) {
+
+            response.addHeader(HEADER_LIMIT_REMAINING, String.valueOf(probe.getRemainingTokens()));
+            return true;
+
+        } else {
+
+            long waitForRefill = probe.getNanosToWaitForRefill() / 1_000_000_000;
+
+            response.setContentType(MediaType.APPLICATION_JSON_VALUE);
+            response.addHeader(HEADER_RETRY_AFTER, String.valueOf(waitForRefill));
+            response.sendError(HttpStatus.TOO_MANY_REQUESTS.value(), "You have exhausted your API Request Quota"); // 429
+
+            return false;
+        }
+    }
+}

spring-boot-modules/spring-boot-libraries/src/main/java/com/baeldung/ratelimiting/bucket4japp/service/PricingPlan.java

@@ -0,0 +1,42 @@
+package com.baeldung.ratelimiting.bucket4japp.service;
+
+import java.time.Duration;
+
+import io.github.bucket4j.Bandwidth;
+import io.github.bucket4j.Refill;
+
+public enum PricingPlan {
+
+    FREE(20),
+
+    BASIC(40),
+
+    PROFESSIONAL(100);
+
+    private int bucketCapacity;
+    
+    private PricingPlan(int bucketCapacity) {
+        this.bucketCapacity = bucketCapacity;
+    }
+    
+    Bandwidth getLimit() {
+        return Bandwidth.classic(bucketCapacity, Refill.intervally(bucketCapacity, Duration.ofHours(1)));
+    }
+    
+    public int bucketCapacity() {
+        return bucketCapacity;
+    }
+    
+    static PricingPlan resolvePlanFromApiKey(String apiKey) {
+        if (apiKey == null || apiKey.isEmpty()) {
+            return FREE;
+        
+        } else if (apiKey.startsWith("PX001-")) {
+            return PROFESSIONAL;
+            
+        } else if (apiKey.startsWith("BX001-")) {
+            return BASIC;
+        }
+        return FREE;
+    }
+}

spring-boot-modules/spring-boot-libraries/src/main/java/com/baeldung/ratelimiting/bucket4japp/service/PricingPlanService.java

@@ -0,0 +1,31 @@
+package com.baeldung.ratelimiting.bucket4japp.service;
+
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+
+import org.springframework.stereotype.Service;
+
+import io.github.bucket4j.Bandwidth;
+import io.github.bucket4j.Bucket;
+import io.github.bucket4j.Bucket4j;
+
+@Service
+public class PricingPlanService {
+
+    private final Map<String, Bucket> cache = new ConcurrentHashMap<>();
+
+    public Bucket resolveBucket(String apiKey) {
+        return cache.computeIfAbsent(apiKey, this::newBucket);
+    }
+
+    private Bucket newBucket(String apiKey) {
+        PricingPlan pricingPlan = PricingPlan.resolvePlanFromApiKey(apiKey);
+        return bucket(pricingPlan.getLimit());
+    }
+
+    private Bucket bucket(Bandwidth limit) {
+        return Bucket4j.builder()
+            .addLimit(limit)
+            .build();
+    }
+}

spring-boot-modules/spring-boot-libraries/src/main/java/com/baeldung/ratelimiting/controller/AreaCalculationController.java

@@ -0,0 +1,29 @@
+package com.baeldung.ratelimiting.controller;
+
+import org.springframework.http.MediaType;
+import org.springframework.http.ResponseEntity;
+import org.springframework.web.bind.annotation.PostMapping;
+import org.springframework.web.bind.annotation.RequestBody;
+import org.springframework.web.bind.annotation.RequestMapping;
+import org.springframework.web.bind.annotation.RestController;
+
+import com.baeldung.ratelimiting.dto.AreaV1;
+import com.baeldung.ratelimiting.dto.RectangleDimensionsV1;
+import com.baeldung.ratelimiting.dto.TriangleDimensionsV1;
+
+@RestController
+@RequestMapping(value = "/api/v1/area", consumes = MediaType.APPLICATION_JSON_VALUE)
+class AreaCalculationController {
+
+    @PostMapping(value = "/rectangle", produces = MediaType.APPLICATION_JSON_VALUE)
+    public ResponseEntity<AreaV1> rectangle(@RequestBody RectangleDimensionsV1 dimensions) {
+
+        return ResponseEntity.ok(new AreaV1("rectangle", dimensions.getLength() * dimensions.getWidth()));
+    }
+
+    @PostMapping(value = "/triangle", produces = MediaType.APPLICATION_JSON_VALUE)
+    public ResponseEntity<AreaV1> triangle(@RequestBody TriangleDimensionsV1 dimensions) {
+
+        return ResponseEntity.ok(new AreaV1("triangle", 0.5d * dimensions.getHeight() * dimensions.getBase()));
+    }
+}

spring-boot-modules/spring-boot-libraries/src/main/java/com/baeldung/ratelimiting/dto/AreaV1.java

@@ -0,0 +1,20 @@
+package com.baeldung.ratelimiting.dto;
+
+public class AreaV1 {
+
+    private String shape;
+    private Double area;
+    
+    public AreaV1(String shape, Double area) {
+        this.area = area;
+        this.shape = shape;
+    }
+    
+    public Double getArea() {
+        return area;
+    }
+    
+    public String getShape() {
+        return shape;
+    }
+}

spring-boot-modules/spring-boot-libraries/src/main/java/com/baeldung/ratelimiting/dto/RectangleDimensionsV1.java

@@ -0,0 +1,15 @@
+package com.baeldung.ratelimiting.dto;
+
+public class RectangleDimensionsV1 {
+
+  private double length;
+  private double width;
+
+  public double getLength() {
+    return length;
+  }
+
+  public double getWidth() {
+    return width;
+  }
+}

spring-boot-modules/spring-boot-libraries/src/main/java/com/baeldung/ratelimiting/dto/TriangleDimensionsV1.java

@@ -0,0 +1,15 @@
+package com.baeldung.ratelimiting.dto;
+
+public class TriangleDimensionsV1 {
+
+    private double base;
+    private double height;
+
+    public double getBase() {
+      return base;
+    }
+
+    public double getHeight() {
+      return height;
+    }
+}

spring-boot-modules/spring-boot-libraries/src/main/resources/ratelimiting/application-bucket4j-starter.yml

@@ -0,0 +1,40 @@
+server:
+  port: 9001
+
+spring:
+  application:
+    name: bucket4j-starter-api-rate-limit-app
+  mvc:
+    throw-exception-if-no-handler-found: true
+  resources:
+    add-mappings: false
+  cache:
+    cache-names:
+    - rate-limit-buckets
+    caffeine:
+      spec: maximumSize=100000,expireAfterAccess=3600s
+
+bucket4j:
+  enabled: true
+  filters:
+  - cache-name: rate-limit-buckets
+    url: /api/v1/area.*
+    http-response-body: "{ \"status\": 429, \"error\": \"Too Many Requests\", \"message\": \"You have exhausted your API Request Quota\" }"
+    rate-limits:
+    - expression: "getHeader('X-api-key')"
+      execute-condition: "getHeader('X-api-key').startsWith('PX001-')"
+      bandwidths:
+      - capacity: 100
+        time: 1
+        unit: hours
+    - expression: "getHeader('X-api-key')"
+      execute-condition: "getHeader('X-api-key').startsWith('BX001-')"
+      bandwidths:
+      - capacity: 40
+        time: 1
+        unit: hours
+    - expression: "getHeader('X-api-key')"
+      bandwidths:
+      - capacity: 20
+        time: 1
+        unit: hours

spring-boot-modules/spring-boot-libraries/src/main/resources/ratelimiting/application-bucket4j.yml

@@ -0,0 +1,10 @@
+server:
+  port: 9000
+
+spring:
+  application:
+    name: bucket4j-api-rate-limit-app
+  mvc:
+    throw-exception-if-no-handler-found: true
+  resources:
+    add-mappings: false

spring-boot-modules/spring-boot-libraries/src/test/java/com/baeldung/ratelimiting/bootstarterapp/Bucket4jBootStarterRateLimitIntegrationTest.java

@@ -0,0 +1,63 @@
+package com.baeldung.ratelimiting.bootstarterapp;
+
+import static org.hamcrest.CoreMatchers.equalTo;
+import static org.springframework.test.web.servlet.request.MockMvcRequestBuilders.post;
+import static org.springframework.test.web.servlet.result.MockMvcResultMatchers.header;
+import static org.springframework.test.web.servlet.result.MockMvcResultMatchers.jsonPath;
+import static org.springframework.test.web.servlet.result.MockMvcResultMatchers.status;
+
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.boot.test.autoconfigure.web.servlet.AutoConfigureMockMvc;
+import org.springframework.boot.test.context.SpringBootTest;
+import org.springframework.http.MediaType;
+import org.springframework.test.context.TestPropertySource;
+import org.springframework.test.context.junit4.SpringRunner;
+import org.springframework.test.web.servlet.MockMvc;
+import org.springframework.test.web.servlet.RequestBuilder;
+
+import com.baeldung.ratelimiting.bucket4japp.service.PricingPlan;
+
+@RunWith(SpringRunner.class)
+@SpringBootTest(classes = Bucket4jRateLimitApp.class)
+@TestPropertySource(properties = "spring.config.location=classpath:ratelimiting/application-bucket4j-starter.yml")
+@AutoConfigureMockMvc
+public class Bucket4jBootStarterRateLimitIntegrationTest {
+
+    @Autowired
+    private MockMvc mockMvc;
+
+    @Test
+    public void givenTriangleAreaCalculator_whenRequestsWithinRateLimit_thenAccepted() throws Exception {
+
+        RequestBuilder request = post("/api/v1/area/triangle").contentType(MediaType.APPLICATION_JSON_VALUE)
+            .content("{ \"height\": 8, \"base\": 10 }")
+            .header("X-api-key", "FX001-UBSZ5YRYQ");
+
+        for (int i = 1; i <= PricingPlan.FREE.bucketCapacity(); i++) {
+            mockMvc.perform(request)
+                .andExpect(status().isOk())
+                .andExpect(header().exists("X-Rate-Limit-Remaining"))
+                .andExpect(jsonPath("$.shape", equalTo("triangle")))
+                .andExpect(jsonPath("$.area", equalTo(40d)));
+        }
+    }
+
+    @Test
+    public void givenTriangleAreaCalculator_whenRequestRateLimitTriggered_thenRejected() throws Exception {
+
+        RequestBuilder request = post("/api/v1/area/triangle").contentType(MediaType.APPLICATION_JSON_VALUE)
+            .content("{ \"height\": 8, \"base\": 10 }")
+            .header("X-api-key", "FX001-ZBSY6YSLP");
+
+        for (int i = 1; i <= PricingPlan.FREE.bucketCapacity(); i++) {
+            mockMvc.perform(request); // exhaust limit
+        }
+
+        mockMvc.perform(request)
+            .andExpect(status().isTooManyRequests())
+            .andExpect(jsonPath("$.message", equalTo("You have exhausted your API Request Quota")))
+            .andExpect(header().exists("X-Rate-Limit-Retry-After-Seconds"));
+    }
+}

spring-boot-modules/spring-boot-libraries/src/test/java/com/baeldung/ratelimiting/bucket4japp/Bucket4jRateLimitIntegrationTest.java

@@ -0,0 +1,61 @@
+package com.baeldung.ratelimiting.bucket4japp;
+
+import static org.hamcrest.CoreMatchers.equalTo;
+import static org.springframework.test.web.servlet.request.MockMvcRequestBuilders.post;
+import static org.springframework.test.web.servlet.result.MockMvcResultMatchers.header;
+import static org.springframework.test.web.servlet.result.MockMvcResultMatchers.jsonPath;
+import static org.springframework.test.web.servlet.result.MockMvcResultMatchers.status;
+
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.boot.test.autoconfigure.web.servlet.AutoConfigureMockMvc;
+import org.springframework.boot.test.context.SpringBootTest;
+import org.springframework.http.MediaType;
+import org.springframework.test.context.junit4.SpringRunner;
+import org.springframework.test.web.servlet.MockMvc;
+import org.springframework.test.web.servlet.RequestBuilder;
+
+import com.baeldung.ratelimiting.bucket4japp.service.PricingPlan;
+
+@RunWith(SpringRunner.class)
+@SpringBootTest(classes = Bucket4jRateLimitApp.class)
+@AutoConfigureMockMvc
+public class Bucket4jRateLimitIntegrationTest {
+
+    @Autowired
+    private MockMvc mockMvc;
+
+    @Test
+    public void givenRectangleAreaCalculator_whenRequestsWithinRateLimit_thenAccepted() throws Exception {
+
+        RequestBuilder request = post("/api/v1/area/rectangle").contentType(MediaType.APPLICATION_JSON_VALUE)
+            .content("{ \"length\": 12, \"width\": 10 }")
+            .header("X-api-key", "FX001-UBSZ5YRYQ");
+
+        for (int i = 1; i <= PricingPlan.FREE.bucketCapacity(); i++) {
+            mockMvc.perform(request)
+                .andExpect(status().isOk())
+                .andExpect(header().exists("X-Rate-Limit-Remaining"))
+                .andExpect(jsonPath("$.shape", equalTo("rectangle")))
+                .andExpect(jsonPath("$.area", equalTo(120d)));
+        }
+    }
+
+    @Test
+    public void givenReactangleAreaCalculator_whenRequestRateLimitTriggered_thenRejected() throws Exception {
+
+        RequestBuilder request = post("/api/v1/area/rectangle").contentType(MediaType.APPLICATION_JSON_VALUE)
+            .content("{ \"length\": 12, \"width\": 10 }")
+            .header("X-api-key", "FX001-ZBSY6YSLP");
+
+        for (int i = 1; i <= PricingPlan.FREE.bucketCapacity(); i++) {
+            mockMvc.perform(request); // exhaust limit
+        }
+
+        mockMvc.perform(request)
+            .andExpect(status().isTooManyRequests())
+            .andExpect(status().reason("You have exhausted your API Request Quota"))
+            .andExpect(header().exists("X-Rate-Limit-Retry-After-Seconds"));
+    }
+}

spring-boot-modules/spring-boot-libraries/src/test/java/com/baeldung/ratelimiting/bucket4japp/Bucket4jUsageUnitTest.java

@@ -0,0 +1,82 @@
+package com.baeldung.ratelimiting.bucket4japp;
+
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+
+import java.time.Duration;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+
+import org.junit.jupiter.api.Test;
+
+import io.github.bucket4j.Bandwidth;
+import io.github.bucket4j.Bucket;
+import io.github.bucket4j.Bucket4j;
+import io.github.bucket4j.Refill;
+
+public class Bucket4jUsageUnitTest {
+
+    @Test
+    public void givenBucketLimit_whenExceedLimit_thenConsumeReturnsFalse() {
+        Refill refill = Refill.intervally(10, Duration.ofMinutes(1));
+        Bandwidth limit = Bandwidth.classic(10, refill);
+        Bucket bucket = Bucket4j.builder()
+            .addLimit(limit)
+            .build();
+
+        for (int i = 1; i <= 10; i++) {
+            assertTrue(bucket.tryConsume(1));
+        }
+        assertFalse(bucket.tryConsume(1));
+    }
+
+    @Test
+    public void givenMultipletLimits_whenExceedSmallerLimit_thenConsumeReturnsFalse() {
+        Bucket bucket = Bucket4j.builder()
+            .addLimit(Bandwidth.classic(10, Refill.intervally(10, Duration.ofMinutes(1))))
+            .addLimit(Bandwidth.classic(5, Refill.intervally(5, Duration.ofSeconds(20))))
+            .build();
+
+        for (int i = 1; i <= 5; i++) {
+            assertTrue(bucket.tryConsume(1));
+        }
+        assertFalse(bucket.tryConsume(1));
+    }
+
+    @Test
+    public void givenBucketLimit_whenThrottleRequests_thenConsumeReturnsTrue() throws InterruptedException {
+        Refill refill = Refill.intervally(1, Duration.ofSeconds(2));
+        Bandwidth limit = Bandwidth.classic(1, refill);
+        Bucket bucket = Bucket4j.builder()
+            .addLimit(limit)
+            .build();
+
+        assertTrue(bucket.tryConsume(1));
+
+        ScheduledExecutorService executor = Executors.newScheduledThreadPool(1);
+        CountDownLatch latch = new CountDownLatch(1);
+
+        executor.schedule(new AssertTryConsume(bucket, latch), 2, TimeUnit.SECONDS);
+
+        latch.await();
+    }
+
+    static class AssertTryConsume implements Runnable {
+
+        private Bucket bucket;
+        private CountDownLatch latch;
+
+        AssertTryConsume(Bucket bucket, CountDownLatch latch) {
+            this.bucket = bucket;
+            this.latch = latch;
+        }
+
+        @Override
+        public void run() {
+            assertTrue(bucket.tryConsume(1));
+            latch.countDown();
+        }
+    }
+}

spring-boot-modules/spring-boot-libraries/src/test/java/com/baeldung/ratelimiting/bucket4japp/PricingPlanServiceUnitTest.java

@@ -0,0 +1,36 @@
+package com.baeldung.ratelimiting.bucket4japp;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+import org.junit.jupiter.api.Test;
+
+import com.baeldung.ratelimiting.bucket4japp.service.PricingPlan;
+import com.baeldung.ratelimiting.bucket4japp.service.PricingPlanService;
+
+import io.github.bucket4j.Bucket;
+
+public class PricingPlanServiceUnitTest {
+
+    private PricingPlanService service = new PricingPlanService();
+
+    @Test
+    public void givenAPIKey_whenFreePlan_thenReturnFreePlanBucket() {
+        Bucket bucket = service.resolveBucket("FX001-UBSZ5YRYQ");
+
+        assertEquals(PricingPlan.FREE.bucketCapacity(), bucket.getAvailableTokens());
+    }
+
+    @Test
+    public void givenAPIKey_whenBasiclan_thenReturnBasicPlanBucket() {
+        Bucket bucket = service.resolveBucket("BX001-MBSZ5YRYP");
+
+        assertEquals(PricingPlan.BASIC.bucketCapacity(), bucket.getAvailableTokens());
+    }
+
+    @Test
+    public void givenAPIKey_whenProfessionalPlan_thenReturnProfessionalPlanBucket() {
+        Bucket bucket = service.resolveBucket("PX001-NBSZ5YRYY");
+
+        assertEquals(PricingPlan.PROFESSIONAL.bucketCapacity(), bucket.getAvailableTokens());
+    }
+}