test(image): add unit tests for _to_image_tensor (#281)

alessiamarcolini · alaeddine-13 · web-flow · commit 6ae22db0c279 · 2022-04-20T02:34:53.000+02:00
* test(image): add unit tests for _to_image_tensor

* style(image): use single-quote strings

* chore: rename variable

Co-authored-by: AlaeddineAbdessalem &lt;alaeddine-13@live.fr&gt;
diff --git a/tests/unit/document/test_image.py b/tests/unit/document/test_image.py
@@ -0,0 +1,42 @@
+import os
+
+import numpy as np
+import pytest
+from PIL import Image
+import io
+
+from docarray.document.mixins.image import _to_image_tensor
+
+
+@pytest.mark.parametrize(
+    'width, height, output_shape',
+    [
+        (None, None, (50, 10, 3)),
+        (8, None, (50, 8, 3)),
+        (None, 30, (30, 10, 3)),
+        (30, 8, (8, 30, 3)),
+    ],
+)
+def test_to_image_tensor_pil(rgb_image_path, width, height, output_shape):
+    tensor = _to_image_tensor(rgb_image_path, width, height)
+
+    assert tensor.shape == output_shape
+    assert isinstance(tensor, np.ndarray)
+
+
+def test_to_image_tensor_blob(rgb_image_path):
+    with open(rgb_image_path, 'rb') as f:
+        blob = io.BytesIO(f.read())
+
+    tensor = _to_image_tensor(blob, None, None)
+
+    assert tensor.shape == (50, 10, 3)
+    assert isinstance(tensor, np.ndarray)
+
+
+@pytest.fixture
+def rgb_image_path(tmpdir):
+    img_path = os.path.join(tmpdir, 'image.png')
+    image = Image.new('RGB', size=(10, 50), color=(0, 0, 0))
+    image.save(img_path)
+    return img_path
