from itertools import starmap, repeat import numpy as np import pytest from PIL import Image from funcy import lmap, compose, flatten from pyinfra.server.utils import pack, unpack, normalize_item from pyinfra.utils.func import star, lift, lstarlift from test.utils.image import image_to_bytes @pytest.fixture def data(data_type, pdf): if data_type == "pdf": return pdf elif data_type == "bytestring": return "content".encode("latin1") @pytest.fixture 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(lift(star(pack)), normalize_item, operation) expected = lmap(unpack, flatten(starmap(op, zip(input_data_items, metadata)))) return expected @pytest.fixture def endpoint(url): return f"{url}/submit" @pytest.fixture(params=[1, 0, 5]) def n_items(request): return request.param @pytest.fixture(params=[0, 100]) def n_pages(request): return request.param @pytest.fixture(params=[1, 5]) def buffer_size(request): return request.param def array_to_image(array) -> Image.Image: return Image.fromarray(np.uint8(array * 255), mode="RGB") def input_batch(n_items): return np.random.random_sample(size=(n_items, 3, 30, 30)) def images(n_items): return lmap(compose(image_to_bytes, array_to_image), input_batch(n_items)) @pytest.fixture def metadata(n_items): return list(repeat({}, n_items)) @pytest.fixture def packages(input_data_items, metadata): return lstarlift(pack)(zip(input_data_items, metadata))