diff --git a/src/apify_client/_resource_clients/request_queue.py b/src/apify_client/_resource_clients/request_queue.py index baa1bf30..c927d773 100644 --- a/src/apify_client/_resource_clients/request_queue.py +++ b/src/apify_client/_resource_clients/request_queue.py @@ -40,12 +40,10 @@ from apify_client._resource_clients._resource_client import ResourceClient, ResourceClientAsync from apify_client._utils.errors import catch_not_found_or_throw from apify_client._utils.http import response_to_dict, to_path_segment -from apify_client._utils.time import to_seconds from apify_client.errors import ApifyApiError if TYPE_CHECKING: from collections.abc import AsyncIterator, Iterator - from datetime import timedelta from apify_client._literals import GeneralAccess from apify_client._typeddicts import ( @@ -197,7 +195,7 @@ def list_head(self, *, limit: int | None = None, timeout: Timeout = 'short') -> def list_and_lock_head( self, *, - lock_duration: timedelta, + lock_secs: int, limit: int | None = None, timeout: Timeout = 'medium', ) -> LockedRequestQueueHead: @@ -206,7 +204,7 @@ def list_and_lock_head( https://docs.apify.com/api/v2#/reference/request-queues/queue-head-with-locks/get-head-and-lock Args: - lock_duration: How long the requests will be locked for. + lock_secs: How long the requests will be locked for, in seconds. limit: How many requests to retrieve. timeout: Timeout for the API HTTP request. @@ -214,7 +212,7 @@ def list_and_lock_head( The desired number of locked requests from the beginning of the queue. """ request_params = self._build_params( - lockSecs=to_seconds(lock_duration, as_int=True), + lockSecs=lock_secs, limit=limit, clientKey=self.client_key, ) @@ -357,7 +355,7 @@ def prolong_request_lock( request_id: str, *, forefront: bool | None = None, - lock_duration: timedelta, + lock_secs: int, timeout: Timeout = 'medium', ) -> RequestLockInfo | None: """Prolong the lock on a request. @@ -367,13 +365,13 @@ def prolong_request_lock( Args: request_id: ID of the request to prolong the lock. forefront: Whether to put the request in the beginning or the end of the queue after lock expires. - lock_duration: By how much to prolong the lock. + lock_secs: By how many seconds to prolong the lock. timeout: Timeout for the API HTTP request. """ request_params = self._build_params( clientKey=self.client_key, forefront=forefront, - lockSecs=to_seconds(lock_duration, as_int=True), + lockSecs=lock_secs, ) response = self._http_client.call( @@ -732,7 +730,7 @@ async def list_head(self, *, limit: int | None = None, timeout: Timeout = 'short async def list_and_lock_head( self, *, - lock_duration: timedelta, + lock_secs: int, limit: int | None = None, timeout: Timeout = 'medium', ) -> LockedRequestQueueHead: @@ -741,7 +739,7 @@ async def list_and_lock_head( https://docs.apify.com/api/v2#/reference/request-queues/queue-head-with-locks/get-head-and-lock Args: - lock_duration: How long the requests will be locked for. + lock_secs: How long the requests will be locked for, in seconds. limit: How many requests to retrieve. timeout: Timeout for the API HTTP request. @@ -749,7 +747,7 @@ async def list_and_lock_head( The desired number of locked requests from the beginning of the queue. """ request_params = self._build_params( - lockSecs=to_seconds(lock_duration, as_int=True), + lockSecs=lock_secs, limit=limit, clientKey=self.client_key, ) @@ -888,7 +886,7 @@ async def prolong_request_lock( request_id: str, *, forefront: bool | None = None, - lock_duration: timedelta, + lock_secs: int, timeout: Timeout = 'medium', ) -> RequestLockInfo | None: """Prolong the lock on a request. @@ -898,13 +896,13 @@ async def prolong_request_lock( Args: request_id: ID of the request to prolong the lock. forefront: Whether to put the request in the beginning or the end of the queue after lock expires. - lock_duration: By how much to prolong the lock. + lock_secs: By how many seconds to prolong the lock. timeout: Timeout for the API HTTP request. """ request_params = self._build_params( clientKey=self.client_key, forefront=forefront, - lockSecs=to_seconds(lock_duration, as_int=True), + lockSecs=lock_secs, ) response = await self._http_client.call( diff --git a/tests/integration/test_request_queue.py b/tests/integration/test_request_queue.py index e2428917..f9ab194c 100644 --- a/tests/integration/test_request_queue.py +++ b/tests/integration/test_request_queue.py @@ -3,7 +3,7 @@ from __future__ import annotations from collections.abc import AsyncIterator, Iterator -from datetime import UTC, datetime, timedelta +from datetime import UTC, datetime from typing import TYPE_CHECKING, Any import pytest @@ -169,7 +169,7 @@ async def test_request_queue_lock(client: ApifyClient | ApifyClientAsync) -> Non # Poll until all requests are available for locking (eventual consistency) async def lock_head() -> LockedRequestQueueHead: - head = await maybe_await(rq.list_and_lock_head(limit=10, lock_duration=timedelta(seconds=10))) + head = await maybe_await(rq.list_and_lock_head(limit=10, lock_secs=10)) assert isinstance(head, LockedRequestQueueHead) return head @@ -200,7 +200,7 @@ async def lock_head() -> LockedRequestQueueHead: prolong_request_lock_response = await maybe_await( rq.prolong_request_lock( get_head_and_lock_response.items[3].id, - lock_duration=timedelta(seconds=15), + lock_secs=15, ) ) assert isinstance(prolong_request_lock_response, RequestLockInfo) @@ -476,7 +476,7 @@ async def test_request_queue_list_and_lock_head(client: ApifyClient | ApifyClien await ensure_queue_is_populated(rq_client, expected_count=5) - result = await maybe_await(rq_client.list_and_lock_head(limit=3, lock_duration=timedelta(seconds=60))) + result = await maybe_await(rq_client.list_and_lock_head(limit=3, lock_secs=60)) assert isinstance(result, LockedRequestQueueHead) lock_response = result assert len(lock_response.items) == 3 @@ -503,7 +503,7 @@ async def test_request_queue_prolong_request_lock(client: ApifyClient | ApifyCli # Poll until the request is available for locking (eventual consistency) async def lock_head() -> LockedRequestQueueHead: - head = await maybe_await(rq_client.list_and_lock_head(limit=1, lock_duration=timedelta(seconds=60))) + head = await maybe_await(rq_client.list_and_lock_head(limit=1, lock_secs=60)) assert isinstance(head, LockedRequestQueueHead) return head @@ -513,9 +513,7 @@ async def lock_head() -> LockedRequestQueueHead: original_lock_expires = locked_request.lock_expires_at # Prolong the lock - prolong_response = await maybe_await( - rq_client.prolong_request_lock(locked_request.id, lock_duration=timedelta(seconds=120)) - ) + prolong_response = await maybe_await(rq_client.prolong_request_lock(locked_request.id, lock_secs=120)) assert isinstance(prolong_response, RequestLockInfo) assert prolong_response.lock_expires_at is not None assert prolong_response.lock_expires_at > original_lock_expires @@ -537,7 +535,7 @@ async def test_request_queue_delete_request_lock(client: ApifyClient | ApifyClie # Poll until the request is available for locking (eventual consistency) async def lock_head() -> LockedRequestQueueHead: - head = await maybe_await(rq_client.list_and_lock_head(limit=1, lock_duration=timedelta(seconds=60))) + head = await maybe_await(rq_client.list_and_lock_head(limit=1, lock_secs=60)) assert isinstance(head, LockedRequestQueueHead) return head @@ -573,7 +571,7 @@ async def test_request_queue_unlock_requests(client: ApifyClient | ApifyClientAs await ensure_queue_is_populated(rq_client, expected_count=5) - lock_response = await maybe_await(rq_client.list_and_lock_head(limit=3, lock_duration=timedelta(seconds=60))) + lock_response = await maybe_await(rq_client.list_and_lock_head(limit=3, lock_secs=60)) assert isinstance(lock_response, LockedRequestQueueHead) assert len(lock_response.items) == 3 diff --git a/tests/unit/test_request_queue_locks.py b/tests/unit/test_request_queue_locks.py new file mode 100644 index 00000000..d7af217a --- /dev/null +++ b/tests/unit/test_request_queue_locks.py @@ -0,0 +1,105 @@ +from __future__ import annotations + +import json +from typing import TYPE_CHECKING + +import pytest +from werkzeug import Request, Response + +from apify_client import ApifyClient, ApifyClientAsync + +if TYPE_CHECKING: + from collections.abc import Callable + + from pytest_httpserver import HTTPServer + +pytestmark = pytest.mark.usefixtures('http_client_classes') + +_QUEUE_ID = 'test-queue-id' +_REQUEST_ID = 'test-request-id' +_HEAD_LOCK_PATH = f'/v2/request-queues/{_QUEUE_ID}/head/lock' +_REQUEST_LOCK_PATH = f'/v2/request-queues/{_QUEUE_ID}/requests/{_REQUEST_ID}/lock' +_HEAD_LOCK_RESPONSE = { + 'data': { + 'limit': 1, + 'queueModifiedAt': '2024-01-01T00:00:00.000Z', + 'hadMultipleClients': False, + 'lockSecs': 60, + 'queueHasLockedRequests': True, + 'clientKey': 'test-client-key', + 'items': [], + } +} +_REQUEST_LOCK_RESPONSE = {'data': {'lockExpiresAt': '2024-01-01T00:01:00.000Z'}} + + +def _capture_into(captured_requests: list[Request], body: dict) -> Callable[[Request], Response]: + def handler(request: Request) -> Response: + captured_requests.append(request) + return Response(status=200, response=json.dumps(body), mimetype='application/json') + + return handler + + +def test_list_and_lock_head_sends_lock_secs_sync(httpserver: HTTPServer) -> None: + """`list_and_lock_head()` sends the lock length unchanged as the `lockSecs` query parameter.""" + captured_requests: list[Request] = [] + httpserver.expect_request(_HEAD_LOCK_PATH, method='POST').respond_with_handler( + _capture_into(captured_requests, _HEAD_LOCK_RESPONSE) + ) + client = ApifyClient(token='test-token', api_url=httpserver.url_for('/').removesuffix('/')) + + head = client.request_queue(_QUEUE_ID).list_and_lock_head(lock_secs=60, limit=1) + + assert head.lock_secs == 60 + assert len(captured_requests) == 1 + assert captured_requests[0].args['lockSecs'] == '60' + assert captured_requests[0].args['limit'] == '1' + + +async def test_list_and_lock_head_sends_lock_secs_async(httpserver: HTTPServer) -> None: + """Async variant of `test_list_and_lock_head_sends_lock_secs_sync`.""" + captured_requests: list[Request] = [] + httpserver.expect_request(_HEAD_LOCK_PATH, method='POST').respond_with_handler( + _capture_into(captured_requests, _HEAD_LOCK_RESPONSE) + ) + client = ApifyClientAsync(token='test-token', api_url=httpserver.url_for('/').removesuffix('/')) + + head = await client.request_queue(_QUEUE_ID).list_and_lock_head(lock_secs=60, limit=1) + + assert head.lock_secs == 60 + assert len(captured_requests) == 1 + assert captured_requests[0].args['lockSecs'] == '60' + assert captured_requests[0].args['limit'] == '1' + + +def test_prolong_request_lock_sends_lock_secs_sync(httpserver: HTTPServer) -> None: + """`prolong_request_lock()` sends the lock length unchanged as the `lockSecs` query parameter.""" + captured_requests: list[Request] = [] + httpserver.expect_request(_REQUEST_LOCK_PATH, method='PUT').respond_with_handler( + _capture_into(captured_requests, _REQUEST_LOCK_RESPONSE) + ) + client = ApifyClient(token='test-token', api_url=httpserver.url_for('/').removesuffix('/')) + + lock_info = client.request_queue(_QUEUE_ID).prolong_request_lock(_REQUEST_ID, lock_secs=120, forefront=True) + + assert lock_info is not None + assert len(captured_requests) == 1 + assert captured_requests[0].args['lockSecs'] == '120' + assert captured_requests[0].args['forefront'] == 'true' + + +async def test_prolong_request_lock_sends_lock_secs_async(httpserver: HTTPServer) -> None: + """Async variant of `test_prolong_request_lock_sends_lock_secs_sync`.""" + captured_requests: list[Request] = [] + httpserver.expect_request(_REQUEST_LOCK_PATH, method='PUT').respond_with_handler( + _capture_into(captured_requests, _REQUEST_LOCK_RESPONSE) + ) + client = ApifyClientAsync(token='test-token', api_url=httpserver.url_for('/').removesuffix('/')) + + lock_info = await client.request_queue(_QUEUE_ID).prolong_request_lock(_REQUEST_ID, lock_secs=120, forefront=True) + + assert lock_info is not None + assert len(captured_requests) == 1 + assert captured_requests[0].args['lockSecs'] == '120' + assert captured_requests[0].args['forefront'] == 'true'