60 lines
1.4 KiB
Python
60 lines
1.4 KiB
Python
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
|
|
|
|
|
|
@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(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
|
|
|
|
|
|
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():
|
|
return {"idx": [1, 100, 101], "path": "asd/asd"}
|