diff --git a/integrations/snyk/snyk/rate_limit.py b/integrations/snyk/snyk/rate_limit.py index fc05ee0fd..9b65612c1 100644 --- a/integrations/snyk/snyk/rate_limit.py +++ b/integrations/snyk/snyk/rate_limit.py @@ -5,7 +5,9 @@ class RateLimiter: def __init__(self, concurrency_limit: int, requests_per_minute: int) -> None: - self.semaphore: asyncio.BoundedSemaphore = asyncio.BoundedSemaphore(concurrency_limit) + self.semaphore: asyncio.BoundedSemaphore = asyncio.BoundedSemaphore( + concurrency_limit + ) self.requests_per_minute: int = requests_per_minute self.request_count: int = 0 self.reset_time: float = time.monotonic() + 60 @@ -19,7 +21,9 @@ async def acquire(self) -> None: self.reset_time = current_time + 60 if self.request_count >= self.requests_per_minute: sleep_time: float = self.reset_time - current_time - logger.info(f"Rate limit reached, sleeping for {sleep_time:.2f} seconds") + logger.info( + f"Rate limit reached, sleeping for {sleep_time:.2f} seconds" + ) await asyncio.sleep(sleep_time) self.request_count = 0 self.reset_time = time.monotonic() + 60 diff --git a/integrations/snyk/tests/snyk/test_rate_limit.py b/integrations/snyk/tests/snyk/test_rate_limit.py index 2401cbe66..734f12035 100644 --- a/integrations/snyk/tests/snyk/test_rate_limit.py +++ b/integrations/snyk/tests/snyk/test_rate_limit.py @@ -4,13 +4,13 @@ from loguru import logger from snyk.rate_limit import RateLimiter + @pytest.mark.asyncio async def test_concurrency_limit() -> None: concurrency_limit: int = 3 requests_per_minute: int = 5 rate_limiter: RateLimiter = RateLimiter( - concurrency_limit=concurrency_limit, - requests_per_minute=requests_per_minute + concurrency_limit=concurrency_limit, requests_per_minute=requests_per_minute ) async def request_task(task_id: int) -> None: @@ -27,6 +27,7 @@ async def request_task(task_id: int) -> None: elapsed_time: float = time.monotonic() - start_time assert elapsed_time >= 1.0 + @pytest.mark.asyncio async def test_rate_limit() -> None: rate_limiter: RateLimiter = RateLimiter(concurrency_limit=5, requests_per_minute=2) @@ -40,6 +41,7 @@ async def make_request() -> None: elapsed_time: float = time.monotonic() - start_time assert elapsed_time >= 60, "Rate limiter did not properly enforce rate limit." + @pytest.mark.asyncio async def test_request_count_reset() -> None: rate_limiter: RateLimiter = RateLimiter(concurrency_limit=2, requests_per_minute=2) @@ -55,5 +57,7 @@ async def test_request_count_reset() -> None: await rate_limiter.acquire() elapsed_time: float = time.monotonic() - start_time - assert elapsed_time < 1, "Rate limiter did not reset request count after 60 seconds." + assert ( + elapsed_time < 1 + ), "Rate limiter did not reset request count after 60 seconds." rate_limiter.release()