supports callable rate limit

AgentSoftware · Apr 5, 2024 · e2afbb1 · e2afbb1
1 parent 09e1af2
commit e2afbb1
Show file tree

Hide file tree

Showing 3 changed files with 47 additions and 4 deletions.
diff --git a/src/RawSqsConnector.php b/src/RawSqsConnector.php
@@ -42,7 +42,8 @@ public function connect(array $config): Queue|RawSqsQueue
         }
 
         if (Arr::get($config, 'rate_limit')) {
-            $rawSqsQueue->setRateLimit($config['rate_limit']);
+            $limit = $config['rate_limit'];
+            $rawSqsQueue->setRateLimit(is_callable($limit) ? $limit() : $limit);
         }
 
         return $rawSqsQueue;

diff --git a/src/RawSqsQueue.php b/src/RawSqsQueue.php
@@ -14,7 +14,7 @@
 class RawSqsQueue extends SqsQueue
 {
     protected string $jobClass;
-    protected ?int $rateLimit = null;
+    protected mixed $rateLimit = null;
 
     public function pop($queue = null): SqsJob|Job|null
     {
@@ -146,10 +146,10 @@ public function setJobClass(string $jobClass): static
     }
 
     /**
-     * @param int $rateLimit
+     * @param int|callable $rateLimit
      * @return $this
      */
-    public function setRateLimit(int $rateLimit): static
+    public function setRateLimit(int|callable $rateLimit): static
     {
         $this->rateLimit = $rateLimit;
         return $this;

diff --git a/tests/RawSqsQueueTest.php b/tests/RawSqsQueueTest.php
@@ -202,6 +202,48 @@ public function testWillReturnMessageIfRateLimitEnabled(): void
         $this->expectNotToPerformAssertions();
     }
 
+    public function testWillReturnMessageIfRateLimitEnabledUsingAClosure(): void
+    {
+        $firstName = 'Primitive';
+        $lastName = 'Sense';
+
+        $sqsReturnMessage = [
+            'Body' => json_encode([
+                'first_name' => $firstName,
+                'last_name' => $lastName
+            ])
+        ];
+
+        $sqsClientMock = Mockery::mock(SqsClient::class);
+        $sqsClientMock->shouldReceive('receiveMessage')
+            ->andReturn([
+                'Messages' => [
+                    $sqsReturnMessage
+                ]
+            ]);
+
+        $rawSqsQueue = Mockery::mock(RawSqsQueue::class, [
+            $sqsClientMock,
+            'default',
+            'prefix'
+        ])
+            ->shouldAllowMockingProtectedMethods()
+            ->makePartial();
+
+        $rawSqsQueue
+            ->shouldReceive('hasRemainingAttempts')
+            ->andReturn(true);
+
+        $container = Mockery::mock(Container::class);
+        $rawSqsQueue->setContainer($container);
+        $rawSqsQueue->setJobClass(TestJobClass::class);
+        $rawSqsQueue->setRateLimit(fn () => 1);
+
+        $rawSqsQueue->pop();
+
+        $this->expectNotToPerformAssertions();
+    }
+
     public function testWillNotReturnMessageIfRateLimitEnabledButNoAttemptsLeft(): void
     {
         $sqsClientMock = Mockery::mock(SqsClient::class);