diff --git a/tests/python/cli/util.py b/tests/python/cli/util.py
index 021d28082f1f..034d5d073ace 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 59a811c7fa8f..4604ab060fcc 100644
--- a/tests/python/sdk/test_jobs.py
+++ b/tests/python/sdk/test_jobs.py
@@ -149,7 +149,12 @@ def test_can_download_frames(self, fxt_new_task: Task, quality: str):
             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 8a5bbf8b5a58..dd5c7b8f2119 100644
--- a/tests/python/sdk/test_tasks.py
+++ b/tests/python/sdk/test_tasks.py
@@ -365,7 +365,12 @@ def test_can_download_frames(self, fxt_new_task: Task, quality: str):
             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 289d24e7966f..312fb99f66c4 100644
--- a/tests/python/shared/utils/helpers.py
+++ b/tests/python/shared/utils/helpers.py
@@ -13,9 +13,9 @@
 
 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