diff --git a/src/apify_client/_resource_clients/run.py b/src/apify_client/_resource_clients/run.py index 563c95ea..8c49d02c 100644 --- a/src/apify_client/_resource_clients/run.py +++ b/src/apify_client/_resource_clients/run.py @@ -390,6 +390,7 @@ def charge( *, count: int = 1, idempotency_key: str | None = None, + request_id: str | None = None, timeout: Timeout = 'short', ) -> None: """Charge for an event of a Pay-Per-Event Actor run. @@ -401,6 +402,9 @@ def charge( count: The number of events to charge. idempotency_key: A unique key to ensure idempotent charging. If not provided, one will be auto-generated. + request_id: ID of the request served by a multi-tenant Actor Server run, taken from the `X-Actor-Request-ID` + header of the incoming request. The charge then goes to the account of the caller who made that request, + not the account that started the run. timeout: Timeout for the API HTTP request. Raises: @@ -414,6 +418,13 @@ def charge( timestamp_ms = int(time.time() * 1000) idempotency_key = f'{self._resource_id}-{event_name}-{timestamp_ms}-{random_suffix}' + body: dict[str, str | int] = { + 'eventName': event_name, + 'count': count, + } + if request_id is not None: + body['requestId'] = request_id + self._http_client.call( url=self._build_url('charge'), method='POST', @@ -421,12 +432,7 @@ def charge( 'idempotency-key': idempotency_key, 'content-type': 'application/json', }, - data=json.dumps( - { - 'eventName': event_name, - 'count': count, - } - ), + data=json.dumps(body), timeout=timeout, ) @@ -940,6 +946,7 @@ async def charge( *, count: int = 1, idempotency_key: str | None = None, + request_id: str | None = None, timeout: Timeout = 'short', ) -> None: """Charge for an event of a Pay-Per-Event Actor run. @@ -951,6 +958,9 @@ async def charge( count: The number of events to charge. idempotency_key: A unique key to ensure idempotent charging. If not provided, one will be auto-generated. + request_id: ID of the request served by a multi-tenant Actor Server run, taken from the `X-Actor-Request-ID` + header of the incoming request. The charge then goes to the account of the caller who made that request, + not the account that started the run. timeout: Timeout for the API HTTP request. Raises: @@ -964,6 +974,13 @@ async def charge( timestamp_ms = int(time.time() * 1000) idempotency_key = f'{self._resource_id}-{event_name}-{timestamp_ms}-{random_suffix}' + body: dict[str, str | int] = { + 'eventName': event_name, + 'count': count, + } + if request_id is not None: + body['requestId'] = request_id + await self._http_client.call( url=self._build_url('charge'), method='POST', @@ -971,12 +988,7 @@ async def charge( 'idempotency-key': idempotency_key, 'content-type': 'application/json', }, - data=json.dumps( - { - 'eventName': event_name, - 'count': count, - } - ), + data=json.dumps(body), timeout=timeout, ) diff --git a/tests/unit/test_run_charge.py b/tests/unit/test_run_charge.py index 069778c6..d062148b 100644 --- a/tests/unit/test_run_charge.py +++ b/tests/unit/test_run_charge.py @@ -71,3 +71,57 @@ def capture_request(request: Request) -> Response: assert len(captured_requests) == 1 body = json.loads(captured_requests[0].get_data()) assert body['count'] == count + + +@pytest.mark.parametrize( + ('request_id', 'expected_body'), + [ + pytest.param(None, {'eventName': 'test-event', 'count': 1}, id='omitted'), + pytest.param('req_123', {'eventName': 'test-event', 'count': 1, 'requestId': 'req_123'}, id='provided'), + ], +) +def test_run_charge_sends_request_id_sync(httpserver: HTTPServer, request_id: str | None, expected_body: dict) -> None: + """Ensure `request_id` is sent as `requestId` in the body, and left out when not provided.""" + captured_requests: list[Request] = [] + + def capture_request(request: Request) -> Response: + captured_requests.append(request) + return Response(status=200, mimetype='application/json') + + httpserver.expect_request(_CHARGE_PATH, method='POST').respond_with_handler(capture_request) + + api_url = httpserver.url_for('/').removesuffix('/') + client = ApifyClient(token='test_token', api_url=api_url) + + client.run(_MOCKED_RUN_ID).charge('test-event', request_id=request_id) + + assert len(captured_requests) == 1 + assert json.loads(captured_requests[0].get_data()) == expected_body + + +@pytest.mark.parametrize( + ('request_id', 'expected_body'), + [ + pytest.param(None, {'eventName': 'test-event', 'count': 1}, id='omitted'), + pytest.param('req_123', {'eventName': 'test-event', 'count': 1, 'requestId': 'req_123'}, id='provided'), + ], +) +async def test_run_charge_sends_request_id_async( + httpserver: HTTPServer, request_id: str | None, expected_body: dict +) -> None: + """Async variant of `test_run_charge_sends_request_id_sync`.""" + captured_requests: list[Request] = [] + + def capture_request(request: Request) -> Response: + captured_requests.append(request) + return Response(status=200, mimetype='application/json') + + httpserver.expect_request(_CHARGE_PATH, method='POST').respond_with_handler(capture_request) + + api_url = httpserver.url_for('/').removesuffix('/') + client = ApifyClientAsync(token='test_token', api_url=api_url) + + await client.run(_MOCKED_RUN_ID).charge('test-event', request_id=request_id) + + assert len(captured_requests) == 1 + assert json.loads(captured_requests[0].get_data()) == expected_body