|
4 | 4 | from __future__ import annotations |
5 | 5 |
|
6 | 6 | import asyncio |
| 7 | +from datetime import timedelta |
7 | 8 | from typing import TYPE_CHECKING |
8 | 9 |
|
9 | 10 | import pytest |
@@ -421,6 +422,36 @@ async def test_is_empty(rq: RequestQueue) -> None: |
421 | 422 | assert await rq.is_empty() is True |
422 | 423 |
|
423 | 424 |
|
| 425 | +@pytest.mark.parametrize( |
| 426 | + ('wait_for_all'), |
| 427 | + [ |
| 428 | + pytest.param(True, id='wait for all'), |
| 429 | + pytest.param(False, id='don`t wait for all'), |
| 430 | + ], |
| 431 | +) |
| 432 | +async def test_add_requests_wait_for_all(rq: RequestQueue, *, wait_for_all: bool) -> None: |
| 433 | + """Test adding requests with wait_for_all_requests_to_be_added option.""" |
| 434 | + urls = [f'https://example.com/{i}' for i in range(15)] |
| 435 | + |
| 436 | + # Add requests without waiting |
| 437 | + await rq.add_requests( |
| 438 | + urls, |
| 439 | + batch_size=5, |
| 440 | + wait_for_all_requests_to_be_added=wait_for_all, |
| 441 | + wait_time_between_batches=timedelta(milliseconds=100), |
| 442 | + ) |
| 443 | + |
| 444 | + if not wait_for_all: |
| 445 | + # Immediately after adding, the total count may be less than 15 due to background processing |
| 446 | + assert await rq.get_total_count() <= 15 |
| 447 | + |
| 448 | + # Wait a 250 milliseconds for background tasks to complete |
| 449 | + await asyncio.sleep(0.25) |
| 450 | + |
| 451 | + # Verify all requests were added |
| 452 | + assert await rq.get_total_count() == 15 |
| 453 | + |
| 454 | + |
424 | 455 | async def test_is_finished(rq: RequestQueue) -> None: |
425 | 456 | """Test checking if a request queue is finished.""" |
426 | 457 | # Initially the queue should be finished (empty and no background tasks) |
|
0 commit comments