From 75e484913e6ad7d6aabcf55ef554fe587121f55a Mon Sep 17 00:00:00 2001
From: Maxim Solovev <m.solovev@paybis.com>
Date: Fri, 14 Jun 2024 15:05:57 +0200
Subject: [PATCH] Add token price and maxTime parameters

---
 src/ThrottlePlugin.php       | 21 ++++++++++++++--
 tests/ThrottlePluginTest.php | 48 ++++++++++++++++++++++++++++++++++++
 2 files changed, 67 insertions(+), 2 deletions(-)

diff --git a/src/ThrottlePlugin.php b/src/ThrottlePlugin.php
index efe0f7b..3ca019f 100644
--- a/src/ThrottlePlugin.php
+++ b/src/ThrottlePlugin.php
@@ -7,20 +7,37 @@
 use Http\Client\Common\Plugin;
 use Http\Promise\Promise;
 use Psr\Http\Message\RequestInterface;
+use Symfony\Component\RateLimiter\Exception\MaxWaitDurationExceededException;
+use Symfony\Component\RateLimiter\Exception\ReserveNotSupportedException;
 use Symfony\Component\RateLimiter\LimiterInterface;
 
 final class ThrottlePlugin implements Plugin
 {
     private LimiterInterface $rateLimiter;
 
-    public function __construct(LimiterInterface $rateLimiter)
+    private int $tokens;
+
+    private ?float $maxTime;
+
+    /**
+     * @param int $tokens the number of tokens required
+     * @param float|null $maxTime maximum accepted waiting time in seconds
+     */
+    public function __construct(LimiterInterface $rateLimiter, int $tokens = 1, ?float $maxTime = null)
     {
         $this->rateLimiter = $rateLimiter;
+        $this->tokens = $tokens;
+        $this->maxTime = $maxTime;
     }
 
+    /**
+     * @throws MaxWaitDurationExceededException if $maxTime is set and the process needs to wait longer than its value (in seconds)
+     * @throws ReserveNotSupportedException     if this limiter implementation doesn't support reserving tokens
+     * @throws \InvalidArgumentException        if $tokens is larger than the maximum burst size
+     */
     public function handleRequest(RequestInterface $request, callable $next, callable $first): Promise
     {
-        $this->rateLimiter->reserve()->wait();
+        $this->rateLimiter->reserve($this->tokens, $this->maxTime)->wait();
 
         return $next($request);
     }
diff --git a/tests/ThrottlePluginTest.php b/tests/ThrottlePluginTest.php
index 1b84db9..a54d9d7 100644
--- a/tests/ThrottlePluginTest.php
+++ b/tests/ThrottlePluginTest.php
@@ -11,6 +11,7 @@
 use Nyholm\Psr7\Request;
 use PHPUnit\Framework\TestCase;
 use Symfony\Bridge\PhpUnit\ClockMock;
+use Symfony\Component\RateLimiter\Exception\MaxWaitDurationExceededException;
 use Symfony\Component\RateLimiter\RateLimit;
 use Symfony\Component\RateLimiter\RateLimiterFactory;
 use Symfony\Component\RateLimiter\Storage\InMemoryStorage;
@@ -56,4 +57,51 @@ public function testThrottle(): void
         $this->client->sendRequest(new Request('GET', ''));
         $this->assertEqualsWithDelta($timeAfterThrottle, time(), 1);
     }
+
+    public function testTokens(): void
+    {
+        $this->client = new PluginClient($this->mockClient, [
+            new ThrottlePlugin(
+                (new RateLimiterFactory(
+                    ['id' => 'foo', 'policy' => 'fixed_window', 'limit' => 2, 'interval' => '3 seconds'],
+                    new InMemoryStorage(),
+                ))->create(),
+                2,
+            ),
+        ]);
+
+        $time = time();
+        $this->client->sendRequest(new Request('GET', ''));
+        $this->client->sendRequest(new Request('GET', ''));
+        $this->assertEqualsWithDelta($time, ($timeAfterThrottle = time()) - 3, 1);
+
+        $this->client->sendRequest(new Request('GET', ''));
+        $this->assertEqualsWithDelta($timeAfterThrottle, time(), 1);
+    }
+
+    public function testMaxTime(): void
+    {
+        $this->client = new PluginClient($this->mockClient, [
+            new ThrottlePlugin(
+                $rateLimit = (new RateLimiterFactory(
+                    ['id' => 'foo', 'policy' => 'fixed_window', 'limit' => 2, 'interval' => '3 seconds'],
+                    new InMemoryStorage(),
+                ))->create(),
+                1,
+                1,
+            ),
+        ]);
+
+        $this->expectException(MaxWaitDurationExceededException::class);
+        $this->expectExceptionMessage('The rate limiter wait time ("3" seconds) is longer than the provided maximum time ("1" seconds).');
+        
+        $time = time();
+        $this->client->sendRequest(new Request('GET', ''));
+        $this->client->sendRequest(new Request('GET', ''));
+        $this->client->sendRequest(new Request('GET', ''));
+        $this->assertEqualsWithDelta($time, ($timeAfterThrottle = time()) - 3, 1);
+
+        $this->client->sendRequest(new Request('GET', ''));
+        $this->assertEqualsWithDelta($timeAfterThrottle, time(), 1);
+    }
 }