add save test for callable data

openvinotoolkit · Nov 13, 2023 · 9dce88f · 9dce88f
1 parent 1a0dcfa
commit 9dce88f
Showing 1 changed file with 13 additions and 0 deletions.
diff --git a/tests/unit/test_image.py b/tests/unit/test_image.py
@@ -5,6 +5,7 @@
 import numpy as np
 
 import datumaro.util.image as image_module
+from datumaro.components.media import ImageFromData
 
 from ..requirements import Requirements, mark_requirement
 
@@ -71,3 +72,15 @@ def test_save_image_can_create_dir(self):
             path = osp.join(test_dir, "some", "path.jpg")
             image_module.save_image(path, np.ones((5, 4, 3)), create_dir=True)
             self.assertTrue(osp.isfile(path))
+
+    @mark_requirement(Requirements.DATUM_GENERAL_REQ)
+    def test_save_callable_image(self):
+        with TestDir() as test_dir:
+            path = osp.join(test_dir, "some", "path.jpg")
+
+            def callable_data():
+                return np.ones((5, 4, 3))
+
+            image = ImageFromData(data=callable_data)
+            image.save(path)
+            self.assertTrue(osp.isfile(path))