Add more tests

src/LeakyBucket.php

@@ -2,6 +2,8 @@
 
 namespace Nadylib\LeakyBucket;
 
+use Closure;
+use Error;
 use Psr\Log\LoggerInterface;
 use Revolt\EventLoop;
 use Revolt\EventLoop\Suspension;
@@ -24,6 +26,30 @@ final public function __construct(
 		$this->fill = $startAmount ?? $this->size;
 	}
 
+	/** Get how much is currently in the bucket */
+	public function getFill(): int {
+		return $this->fill;
+	}
+
+	/** Check if we could take $amount without delay */
+	public function canTake(int $amount=1): bool {
+		return $this->fill < $amount;
+	}
+
+	/**
+	 * Try to take out something from the bucket
+	 *
+	 * If there is currently at least $amount in the bucket, immediately return,
+	 * optionally calling $callback if set. If there is not enough in the bucket,
+	 * wait until there is (if $callback is not set), or immediately return, and
+	 * call $callback once there is enough.
+	 *
+	 * @param int          $amount   How much to take
+	 * @param null|Closure $callback If set, a Closure to call back once there
+	 *                               is enough in the bucket and it's our turn
+	 *
+	 * @throws Error
+	 */
 	public function take(int $amount=1, ?\Closure $callback=null): void {
 		if ($this->fill < $amount) {
 			$this->logger?->debug("Client wants {amount}, bucket has {fill}/{size}, waiting for refill", [
@@ -65,9 +91,7 @@ protected function start(): void {
 		$this->logger?->debug("Starting refill-thread");
 		$this->cancellationToken = EventLoop::repeat(
 			$this->refillDelay,
-			function (string $token): void {
-				$this->put($this->refillAmount);
-			}
+			fn () => $this->put($this->refillAmount)
 		);
 	}
 

tests/BucketTest.php

@@ -7,6 +7,17 @@
 use Revolt\EventLoop;
 
 final class BucketTest extends TestCase {
+	public function testTakingReducesFill(): void {
+		$bucket = new LeakyBucket(
+			size: 5,
+			refillDelay: 0.01,
+		);
+		$bucket->take();
+		$this->assertSame(4, $bucket->getFill());
+		$bucket->take(2);
+		$this->assertSame(2, $bucket->getFill());
+	}
+
 	public function testNoDelayOnFull(): void {
 		$bucket = new LeakyBucket(
 			size: 5,
@@ -57,16 +68,40 @@ public function testPreciseDelayOnRefillWithCallback(): void {
 			refillDelay: 0.2,
 		);
 		$callbacks = 0;
-		$start = microtime(true);
-		for ($i = 1; $i <= 7; $i++) {
-			$bucket->take(callback: function () use (&$callbacks) {
+		$times = [];
+		for ($i = 1; $i <= 8; $i++) {
+			$start = microtime(true);
+			$bucket->take(callback: function () use (&$callbacks, &$times, $start) {
 				$callbacks++;
+				$end = microtime(true);
+				$times []= round($end-$start, 1);
 			});
 		}
 		EventLoop::run();
-		$end = microtime(true);
-		$this->assertSame(7, $callbacks);
-		$this->assertLessThan(1.5, $end-$start);
-		$this->assertGreaterThanOrEqual(1.4, $end-$start);
+		$this->assertSame(8, $callbacks);
+		$this->assertSame($times, [0.0, 0.0, 0.0, 0.0, 0.0, 0.2, 0.4, 0.6]);
+	}
+
+	public function testSleepingRefillsBucket(): void {
+		$bucket = new LeakyBucket(
+			size: 5,
+			refillDelay: 0.1,
+		);
+		$callbacks = 0;
+		$times = [];
+		for ($i = 1; $i <= 5; $i++) {
+			$start = microtime(true);
+			$bucket->take(callback: function () use (&$callbacks, &$times, $start) {
+				$callbacks++;
+				$end = microtime(true);
+				$times []= round($end-$start, 1);
+			});
+		}
+		$this->assertSame(0, $bucket->getFill());
+		$suspension = EventLoop::getSuspension();
+		EventLoop::delay(0.32, fn () => $suspension->resume());
+		$suspension->suspend();
+		$this->assertSame(3, $bucket->getFill());
+		EventLoop::run();
 	}
 }