refactoring

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

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

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):