diff --git a/src/apify_client/_utils/encoding.py b/src/apify_client/_utils/encoding.py index b4adf385..0c5722f0 100644 --- a/src/apify_client/_utils/encoding.py +++ b/src/apify_client/_utils/encoding.py @@ -1,9 +1,9 @@ from __future__ import annotations -import io import json from base64 import b64encode from functools import cache +from inspect import isawaitable, iscoroutine from typing import TYPE_CHECKING, Any from apify_client._models import WebhookCreate, WebhookRepresentation @@ -12,29 +12,51 @@ from apify_client.types import WebhooksList -def encode_key_value_store_record_value(value: Any, *, content_type: str | None = None) -> tuple[Any, str]: +def encode_key_value_store_record_value( + value: Any, *, content_type: str | None = None +) -> tuple[bytes | bytearray | str, str]: """Encode a value for storage in a key-value store record. Args: - value: The value to encode (can be dict, str, bytes, or file-like object). + value: The value to encode. Anything exposing a callable `read` is treated as a file-like object: `read` + is called with no arguments, so the value is consumed from its current position and buffered in + memory whole - the object is neither rewound nor closed, and async file-like objects are rejected. + Any other value is JSON-serialized unless it is already bytes or a string. content_type: The content type; if None, it's inferred from the value type. Returns: A tuple of (encoded_value, content_type). + + Raises: + TypeError: If the value cannot be encoded into a body the transport accepts. """ + # Read file-like values into memory; the transport only accepts bytes-like bodies. Detect them by a + # callable `read` (not `io.IOBase`) so duck-typed file-likes are read, not JSON-serialized. Impit exposes + # no streaming `content=` API, so the value has to be buffered whole. + read = getattr(value, 'read', None) + if callable(read): + value = read() + + if isawaitable(value): + if iscoroutine(value): + value.close() # Prevent a "coroutine was never awaited" warning. + raise TypeError( + 'Async file-like objects are not supported. Await the read yourself and pass the resulting ' + 'bytes or string.' + ) + + if not isinstance(value, (bytes, bytearray, str)): + raise TypeError(f'Reading the file-like value returned {type(value).__name__}, expected bytes or str.') + if not content_type: - if isinstance(value, (bytes, bytearray, io.IOBase)): + if isinstance(value, (bytes, bytearray)): content_type = 'application/octet-stream' elif isinstance(value, str): content_type = 'text/plain; charset=utf-8' else: content_type = 'application/json; charset=utf-8' - if ( - 'application/json' in content_type - and not isinstance(value, (bytes, bytearray, io.IOBase)) - and not isinstance(value, str) - ): + if 'application/json' in content_type and not isinstance(value, (bytes, bytearray, str)): # Don't use indentation to reduce size. value = json.dumps( value, @@ -43,6 +65,14 @@ def encode_key_value_store_record_value(value: Any, *, content_type: str | None default=str, ).encode('utf-8') + # A non-JSON content type skips the serialization above, so anything that is not bytes-like would reach the + # transport unencoded and fail there with an opaque error. + if not isinstance(value, (bytes, bytearray, str)): + raise TypeError( + f'Cannot encode a {type(value).__name__} value as {content_type!r}. Pass bytes, a string, or a ' + 'file-like object, or use a JSON content type.' + ) + return (value, content_type) diff --git a/tests/unit/test_key_value_store.py b/tests/unit/test_key_value_store.py new file mode 100644 index 00000000..eb2176ef --- /dev/null +++ b/tests/unit/test_key_value_store.py @@ -0,0 +1,122 @@ +from __future__ import annotations + +import gzip +import io +from typing import TYPE_CHECKING, Any + +import brotli +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 + + from apify_client.types import HttpCompressionAlgorithm + +_MOCKED_KVS_ID = 'test_kvs_id' +_RECORD_PATH = f'/v2/key-value-stores/{_MOCKED_KVS_ID}/records/f' + + +class DuckTypedReader: + """A file-like object that is not an `io.IOBase`, so only duck-typed detection picks it up.""" + + def read(self) -> bytes: + return b'buffer data' + + +# The values are built by a factory because reading consumes them, and each case runs once per compression +# algorithm. Each case is (value factory, expected uploaded body, expected content type). +_FILE_LIKE_VALUE_CASES = [ + pytest.param(lambda: io.BytesIO(b'buffer data'), b'buffer data', 'application/octet-stream', id='bytes io'), + pytest.param(lambda: io.StringIO('buffer data'), b'buffer data', 'text/plain; charset=utf-8', id='string io'), + pytest.param(DuckTypedReader, b'buffer data', 'application/octet-stream', id='duck-typed reader'), +] + + +@pytest.fixture( + params=[ + pytest.param(('gzip', 'gzip'), id='gzip'), + pytest.param(('brotli', 'br'), id='brotli'), + ] +) +def compression_case(request: pytest.FixtureRequest) -> tuple[HttpCompressionAlgorithm, str]: + """Run each test over both supported request-body compression algorithms, as (algorithm, content-encoding).""" + return request.param + + +@pytest.fixture +def api_url(httpserver: HTTPServer) -> str: + """The base URL of the mock server, in the form the clients expect.""" + return httpserver.url_for('/').removesuffix('/') + + +@pytest.fixture +def captured_records(httpserver: HTTPServer) -> list[Request]: + """Answer record uploads with a 201 and collect the requests the client sent.""" + requests: list[Request] = [] + + def capture_request(request: Request) -> Response: + requests.append(request) + return Response(status=201) + + httpserver.expect_request(_RECORD_PATH, method='PUT').respond_with_handler(capture_request) + return requests + + +def decode_body(request: Request) -> bytes: + """Decompress a captured request body according to its `Content-Encoding`.""" + raw = request.get_data() + encoding = request.headers.get('Content-Encoding') + if encoding == 'gzip': + return gzip.decompress(raw) + if encoding == 'br': + return brotli.decompress(raw) + return raw + + +@pytest.mark.parametrize(('make_value', 'expected_body', 'expected_content_type'), _FILE_LIKE_VALUE_CASES) +def test_set_record_reads_file_like_value_sync( + *, + api_url: str, + captured_records: list[Request], + compression_case: tuple[HttpCompressionAlgorithm, str], + make_value: Callable[[], Any], + expected_body: bytes, + expected_content_type: str, +) -> None: + """A file-like value is read and its bytes are uploaded, not passed through unread.""" + algorithm, content_encoding = compression_case + client = ApifyClient(token='test_token', api_url=api_url, compression=algorithm) + + client.key_value_store(_MOCKED_KVS_ID).set_record('f', make_value()) + + assert len(captured_records) == 1 + assert captured_records[0].headers['content-encoding'] == content_encoding + assert decode_body(captured_records[0]) == expected_body + assert captured_records[0].headers['content-type'] == expected_content_type + + +@pytest.mark.parametrize(('make_value', 'expected_body', 'expected_content_type'), _FILE_LIKE_VALUE_CASES) +async def test_set_record_reads_file_like_value_async( + *, + api_url: str, + captured_records: list[Request], + compression_case: tuple[HttpCompressionAlgorithm, str], + make_value: Callable[[], Any], + expected_body: bytes, + expected_content_type: str, +) -> None: + """A file-like value is read and its bytes are uploaded, not passed through unread.""" + algorithm, content_encoding = compression_case + client = ApifyClientAsync(token='test_token', api_url=api_url, compression=algorithm) + + await client.key_value_store(_MOCKED_KVS_ID).set_record('f', make_value()) + + assert len(captured_records) == 1 + assert captured_records[0].headers['content-encoding'] == content_encoding + assert decode_body(captured_records[0]) == expected_body + assert captured_records[0].headers['content-type'] == expected_content_type diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index a52d3fb4..0ed89bd0 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -223,8 +223,7 @@ def test__clean_json_payload(input_dict: dict, expected: dict) -> None: def test_encode_key_value_store_record_value_dict() -> None: """Test that dictionaries are encoded as JSON.""" value, content_type = encode_key_value_store_record_value({'key': 'value'}) - assert b'"key"' in value - assert b'"value"' in value + assert value == b'{"key": "value"}' assert content_type == 'application/json; charset=utf-8' @@ -249,13 +248,61 @@ def test_encode_key_value_store_record_value( def test_encode_key_value_store_record_value_bytesio() -> None: - """Test that BytesIO is encoded as octet-stream.""" + """Test that BytesIO is read into bytes and encoded as octet-stream.""" buffer = io.BytesIO(b'buffer data') value, content_type = encode_key_value_store_record_value(buffer) - assert value == buffer + assert value == b'buffer data' assert content_type == 'application/octet-stream' +def test_encode_key_value_store_record_value_stringio() -> None: + """Test that StringIO is read into text and encoded as text/plain.""" + buffer = io.StringIO('buffer data') + value, content_type = encode_key_value_store_record_value(buffer) + assert value == 'buffer data' + assert content_type == 'text/plain; charset=utf-8' + + +def test_encode_key_value_store_record_value_duck_typed_file_like() -> None: + """Test that a duck-typed file-like value (a callable `read`, not an `io.IOBase`) is read into bytes.""" + + class Reader: + def read(self) -> bytes: + return b'buffer data' + + value, content_type = encode_key_value_store_record_value(Reader()) + assert value == b'buffer data' + assert content_type == 'application/octet-stream' + + +def test_encode_key_value_store_record_value_async_file_like_raises() -> None: + """Test that an async file-like value is rejected instead of storing the repr of an un-awaited coroutine.""" + + class AsyncReader: + async def read(self) -> bytes: + return b'buffer data' + + with pytest.raises(TypeError, match='Async file-like objects are not supported'): + encode_key_value_store_record_value(AsyncReader()) + + +def test_encode_key_value_store_record_value_non_bytes_read_raises() -> None: + """Test that a `read` returning neither bytes nor str is rejected instead of being JSON-serialized.""" + + class EmptyNonBlockingReader: + def read(self) -> None: + """Mimic a non-blocking raw stream with no data available.""" + + with pytest.raises(TypeError, match='returned NoneType, expected bytes or str'): + encode_key_value_store_record_value(EmptyNonBlockingReader()) + + +def test_encode_key_value_store_record_value_non_encodable_with_explicit_content_type_raises() -> None: + """Test that a non-bytes-like value with a non-JSON content type is rejected before it reaches the transport.""" + with pytest.raises(TypeError, match="Cannot encode a dict value as 'image/png'"): + encode_key_value_store_record_value({'a': 1}, content_type='image/png') + + def test_response_to_dict() -> None: """Test parsing response as dictionary.""" mock_response = Mock()