diff --git a/src/main/kotlin/com/sletmoe/bucket4k/SuspendingBucketConfiguration.kt b/src/main/kotlin/com/sletmoe/bucket4k/SuspendingBucketConfiguration.kt index d027a43..fefee68 100644 --- 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 = mutableListOf(), + private val mutableLimits: MutableList = mutableListOf(), ) { internal val limits: List 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 index b9b9ed7..ab4530f 100644 --- 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) } } }