From 1bc77671e3796baa64944614bc3e4f0a4ffc5516 Mon Sep 17 00:00:00 2001
From: Roman Donchenko <roman@cvat.ai>
Date: Sun, 23 Jul 2023 22:59:28 +0300
Subject: [PATCH] generate_image_file: don't hardcode the file format (#6538)

It doesn't make much sense to encode the image as JPEG if the given file
name is "image.png". Instead, let PIL select the format based on the
requested file name.
---
 tests/python/cli/util.py             | 2 +-
 tests/python/sdk/test_jobs.py        | 7 ++++++-
 tests/python/sdk/test_tasks.py       | 7 ++++++-
 tests/python/shared/utils/helpers.py | 4 ++--
 4 files changed, 15 insertions(+), 5 deletions(-)

diff --git a/tests/python/cli/util.py b/tests/python/cli/util.py
index 021d28082..034d5d073 100644
--- a/tests/python/cli/util.py
+++ b/tests/python/cli/util.py
@@ -26,6 +26,6 @@ def generate_images(dst_dir: Path, count: int) -> List[Path]:
     dst_dir.mkdir(parents=True, exist_ok=True)
     for i in range(count):
         filename = dst_dir / f"img_{i}.jpg"
-        filename.write_bytes(generate_image_file().getvalue())
+        filename.write_bytes(generate_image_file(filename.name).getvalue())
         filenames.append(filename)
     return filenames
diff --git a/tests/python/sdk/test_jobs.py b/tests/python/sdk/test_jobs.py
index 59a811c7f..4604ab060 100644
--- a/tests/python/sdk/test_jobs.py
+++ b/tests/python/sdk/test_jobs.py
@@ -149,7 +149,12 @@ class TestJobUsecases:
             filename_pattern="frame-{frame_id}{frame_ext}",
         )
 
-        assert (self.tmp_path / "frame-0.jpg").is_file()
+        if quality == "original":
+            expected_frame_ext = "png"
+        else:
+            expected_frame_ext = "jpg"
+
+        assert (self.tmp_path / f"frame-0.{expected_frame_ext}").is_file()
         assert self.stdout.getvalue() == ""
 
     def test_can_upload_annotations(self, fxt_new_task: Task, fxt_coco_file: Path):
diff --git a/tests/python/sdk/test_tasks.py b/tests/python/sdk/test_tasks.py
index 8a5bbf8b5..dd5c7b8f2 100644
--- a/tests/python/sdk/test_tasks.py
+++ b/tests/python/sdk/test_tasks.py
@@ -365,7 +365,12 @@ class TestTaskUsecases:
             filename_pattern="frame-{frame_id}{frame_ext}",
         )
 
-        assert (self.tmp_path / "frame-0.jpg").is_file()
+        if quality == "original":
+            expected_frame_ext = "png"
+        else:
+            expected_frame_ext = "jpg"
+
+        assert (self.tmp_path / f"frame-0.{expected_frame_ext}").is_file()
         assert self.stdout.getvalue() == ""
 
     @pytest.mark.parametrize("quality", ("compressed", "original"))
diff --git a/tests/python/shared/utils/helpers.py b/tests/python/shared/utils/helpers.py
index 289d24e79..312fb99f6 100644
--- a/tests/python/shared/utils/helpers.py
+++ b/tests/python/shared/utils/helpers.py
@@ -13,9 +13,9 @@ from shared.fixtures.init import get_server_image_tag
 
 def generate_image_file(filename="image.png", size=(50, 50), color=(0, 0, 0)):
     f = BytesIO()
-    image = Image.new("RGB", size=size, color=color)
-    image.save(f, "jpeg")
     f.name = filename
+    image = Image.new("RGB", size=size, color=color)
+    image.save(f)
     f.seek(0)
 
     return f
-- 
GitLab