refactoring: pulled core operation taking only data out from operation taking data and metadata

test/fixtures/server.py

@@ -1,6 +1,7 @@
 import io
 import socket
 from multiprocessing import Process
+from typing import Generator
 
 import fitz
 import pytest
@@ -10,7 +11,11 @@ from funcy import retry, first, compose
 from waitress import serve
 
 from pyinfra.server.bufferizer.lazy_bufferizer import FlatStreamBuffer
-from pyinfra.server.server import set_up_processing_server, make_streamable_and_wrap_in_packing_logic, QueuedStreamFunction
+from pyinfra.server.server import (
+    set_up_processing_server,
+    make_streamable_and_wrap_in_packing_logic,
+    QueuedStreamFunction,
+)
 from pyinfra.utils.func import starlift
 from test.utils.image import image_to_bytes
 
@@ -54,22 +59,33 @@ def operation_conditionally_batched(operation, batched):
 
 
 @pytest.fixture
-def operation(item_type, one_to_many):
+def operation(core_operation):
-    def upper(string: bytes, metadata):
+    def op(data, metadata):
-        return string.decode().upper().encode(), metadata
+        result = core_operation(data)
+        if isinstance(result, Generator):
+            return zip(result, metadata)
+        else:
+            return result, metadata
+    return op
 
-    def duplicate(string: bytes, metadata):
+
+@pytest.fixture
+def core_operation(item_type, one_to_many):
+    def upper(string: bytes):
+        yield string.decode().upper().encode()
+
+    def duplicate(string: bytes):
         for _ in range(2):
-            yield upper(string, metadata)
+            yield upper(string)
 
-    def rotate(im: bytes, metadata):
+    def rotate(im: bytes):
         im = Image.open(io.BytesIO(im))
-        return image_to_bytes(im.rotate(90)), metadata
+        return image_to_bytes(im.rotate(90))
 
-    def stream_pages(pdf: bytes, metadata):
+    def stream_pages(pdf: bytes):
         for i, page in enumerate(fitz.open(stream=pdf)):
             # yield page.get_pixmap().tobytes("png"), metadata
-            yield f"page_{i}".encode(), metadata
+            yield f"page_{i}".encode()
 
     try:
         return {"string": duplicate if one_to_many else upper, "image": rotate, "pdf": stream_pages}[item_type]