Skip to content

Commit

Permalink
Lint fix
Browse files Browse the repository at this point in the history
  • Loading branch information
oiadebayo committed Oct 22, 2024
1 parent 737d5e2 commit 39cfe2d
Show file tree
Hide file tree
Showing 2 changed files with 13 additions and 5 deletions.
8 changes: 6 additions & 2 deletions integrations/snyk/snyk/rate_limit.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand Down
10 changes: 7 additions & 3 deletions integrations/snyk/tests/snyk/test_rate_limit.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -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)
Expand All @@ -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)
Expand All @@ -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()

0 comments on commit 39cfe2d

Please sign in to comment.