add support for bandwidth builder

ksletmoe · Mar 5, 2024 · 817c040 · 817c040
1 parent 432fb33
commit 817c040
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 6 deletions.
diff --git a/src/main/kotlin/com/sletmoe/bucket4k/SuspendingBucketConfiguration.kt b/src/main/kotlin/com/sletmoe/bucket4k/SuspendingBucketConfiguration.kt
@@ -1,6 +1,8 @@
 package com.sletmoe.bucket4k
 
 import io.github.bucket4j.Bandwidth
+import io.github.bucket4j.BandwidthBuilder.BandwidthBuilderBuildStage
+import io.github.bucket4j.BandwidthBuilder.BandwidthBuilderCapacityStage
 import io.github.bucket4j.TimeMeter
 
 /**
@@ -10,11 +12,11 @@ data class SuspendingBucketConfiguration(
     /**
      * Set's the TimeMeter to be used by the underlying Bucket4j objects. By default, you may choose either
      * TimeMeter.SYSTEM_MILLISECONDS or TimeMeter.SYSTEM_NANOSECONDS. Caution: see Bucket4j's
-     * [documentation about using nanoseconds](https://bucket4j.com/8.2.0/toc.html#customizing-time-measurement-choosing-nanotime-time-resolution)
+     * [documentation about using nanoseconds](https://bucket4j.com/8.9.0/toc.html#customizing-time-measurement-choosing-nanotime-time-resolution)
      * before choosing that option.
      */
     var timeMeter: TimeMeter = TimeMeter.SYSTEM_MILLISECONDS,
-    private var mutableLimits: MutableList<Bandwidth> = mutableListOf(),
+    private val mutableLimits: MutableList<Bandwidth> = mutableListOf(),
 ) {
     internal val limits: List<Bandwidth>
         get() = mutableLimits.toList()
@@ -23,7 +25,16 @@ data class SuspendingBucketConfiguration(
      * Adds a [bandwidth] limit to the bucket. See the Bucket4j documentation on [Bandwidths](https://bucket4j.com/8.2.0/toc.html#bandwidth)
      * for more information.
      */
+    @Deprecated("This method is deprecated, you should use `addLimit { ... }`")
     fun addLimit(bandwidth: Bandwidth) {
         mutableLimits.add(bandwidth)
     }
+
+    /**
+     * Adds a [bandwidth] limit to the bucket. See the Bucket4j documentation on [Bandwidths](https://bucket4j.com/8.9.0/toc.html#bandwidth)
+     * for more information.
+     */
+    fun addLimit(limit: BandwidthBuilderCapacityStage.() -> BandwidthBuilderBuildStage) {
+        mutableLimits.add(Bandwidth.builder().limit().build())
+    }
 }
diff --git a/src/test/kotlin/com/sletmoe/bucket4k/SuspendingBucketTest.kt b/src/test/kotlin/com/sletmoe/bucket4k/SuspendingBucketTest.kt
@@ -1,9 +1,7 @@
 package com.sletmoe.bucket4k
 
-import io.github.bucket4j.Bandwidth
-import io.github.bucket4j.Refill
+import io.kotest.assertions.nondeterministic.eventually
 import io.kotest.assertions.throwables.shouldThrow
-import io.kotest.assertions.timing.eventually
 import io.kotest.core.spec.style.FunSpec
 import io.kotest.core.test.testCoroutineScheduler
 import io.kotest.matchers.shouldBe
@@ -23,7 +21,7 @@ class SuspendingBucketTest : FunSpec() {
         beforeEach {
             // 5 tokens up front (with a max capacity of 5), add one token per 1 second interval
             bucket = SuspendingBucket.build {
-                addLimit(Bandwidth.classic(5, Refill.intervally(1, 1.seconds.toJavaDuration())))
+                addLimit { capacity(5).refillIntervally(1, 1.seconds.toJavaDuration()).initialTokens(5) }
             }
         }