diff --git a/test/exploration_tests/partial_response_test.py b/test/exploration_tests/partial_response_test.py index c93d9bb..defd2d8 100644 --- a/test/exploration_tests/partial_response_test.py +++ b/test/exploration_tests/partial_response_test.py @@ -42,10 +42,14 @@ def post_partial(url, input_data: Iterable[bytes], metadata): return input_data_to_result_data(input_data) -@pytest.mark.parametrize("item_type", ["string"]) -def test_sending_partial_request(url, data_items, metadata, operation, server_process): - op = compose(star(pack), partial(operation, metadata=metadata)) - expected = lmap(op, data_items) - output = list(post_partial(f"{url}/process", data_items, metadata)) +@pytest.mark.parametrize( + "item_type", + [ + "string", + "image", + ], +) +def test_sending_partial_request(url, input_data_items, metadata, operation, target_data_items, server_process): - assert output == expected + output = list(post_partial(f"{url}/process", input_data_items, metadata)) + assert output == target_data_items diff --git a/test/fixtures/input.py b/test/fixtures/input.py index d3a0f78..b583cd6 100644 --- a/test/fixtures/input.py +++ b/test/fixtures/input.py @@ -1,8 +1,12 @@ +from functools import partial + import numpy as np import pytest from PIL import Image from funcy import lmap, compose +from pyinfra.rest import pack +from pyinfra.utils.func import star from test.utils.image import image_to_bytes @@ -15,15 +19,24 @@ def data(data_type, pdf): @pytest.fixture -def data_items(item_type, n_items): +def input_data_items(item_type, n_items, pdf): if item_type == "string": return [bytes(f"content{i}", encoding="utf8") for i in range(n_items)] elif item_type == "image": return images(n_items) + elif item_type == "pdf": + return [pdf] * n_items else: raise ValueError(f"Unknown item type {item_type}") +@pytest.fixture +def target_data_items(input_data_items, item_type, operation, metadata): + op = compose(star(pack), partial(operation, metadata=metadata)) + expected = lmap(op, input_data_items) + return expected + + @pytest.fixture(params=[0, 1, 2, 5]) def n_items(request): return request.param diff --git a/test/fixtures/server.py b/test/fixtures/server.py index f3245c7..e0a94a7 100644 --- a/test/fixtures/server.py +++ b/test/fixtures/server.py @@ -62,7 +62,7 @@ def operation(item_type, batched): def rotate(im: bytes, metadata): im = Image.open(io.BytesIO(im)) - return [(image_to_bytes(im.rotate(90)), metadata)] + return image_to_bytes(im.rotate(90)), metadata def stream_pages(pdf: bytes, metadata): for page in fitz.open(stream=pdf):