From 78e2b726fb597aa6afa87aef70a069568e792e8b Mon Sep 17 00:00:00 2001
From: filipstrand <strand.filip@gmail.com>
Date: Wed, 9 Oct 2024 23:38:56 +0200
Subject: [PATCH] Make sure cleanup happens even if tests fail

---
 ...image_generation_controlnet_test_helper.py | 64 ++++++++++---------
 tests/helpers/image_generation_test_helper.py | 56 ++++++++--------
 2 files changed, 62 insertions(+), 58 deletions(-)

diff --git a/tests/helpers/image_generation_controlnet_test_helper.py b/tests/helpers/image_generation_controlnet_test_helper.py
index c629424..5ed4f6a 100644
--- a/tests/helpers/image_generation_controlnet_test_helper.py
+++ b/tests/helpers/image_generation_controlnet_test_helper.py
@@ -27,37 +27,39 @@ class ImageGeneratorControlnetTestHelper:
         controlnet_image_path = str(ImageGeneratorTestHelper.resolve_path(controlnet_image_path))
         lora_paths = [str(ImageGeneratorTestHelper.resolve_path(p)) for p in lora_paths] if lora_paths else None
 
-        # given
-        flux = Flux1Controlnet(
-            model_config=model_config,
-            quantize=8,
-            lora_paths=lora_paths,
-            lora_scales=lora_scales,
-        )
+        try:
+            # given
+            flux = Flux1Controlnet(
+                model_config=model_config,
+                quantize=8,
+                lora_paths=lora_paths,
+                lora_scales=lora_scales,
+            )
 
-        # when
-        image = flux.generate_image(
-            seed=seed,
-            prompt=prompt,
-            output=str(output_image_path),
-            controlnet_image_path=controlnet_image_path,
-            controlnet_save_canny=False,
-            config=ConfigControlnet(
-                num_inference_steps=steps,
-                height=768,
-                width=493,
-                controlnet_strength=controlnet_strength,
-            ),
-        )
-        image.save(path=output_image_path)
+            # when
+            image = flux.generate_image(
+                seed=seed,
+                prompt=prompt,
+                output=str(output_image_path),
+                controlnet_image_path=controlnet_image_path,
+                controlnet_save_canny=False,
+                config=ConfigControlnet(
+                    num_inference_steps=steps,
+                    height=768,
+                    width=493,
+                    controlnet_strength=controlnet_strength,
+                ),
+            )
+            image.save(path=output_image_path)
 
-        # then
-        np.testing.assert_array_equal(
-            np.array(Image.open(output_image_path)),
-            np.array(Image.open(reference_image_path)),
-            err_msg="Generated image doesn't match reference image",
-        )
+            # then
+            np.testing.assert_array_equal(
+                np.array(Image.open(output_image_path)),
+                np.array(Image.open(reference_image_path)),
+                err_msg="Generated image doesn't match reference image",
+            )
 
-        # cleanup
-        if os.path.exists(output_image_path):
-            os.remove(output_image_path)
+        finally:
+            # cleanup
+            if os.path.exists(output_image_path):
+                os.remove(output_image_path)
diff --git a/tests/helpers/image_generation_test_helper.py b/tests/helpers/image_generation_test_helper.py
index 2c5666b..72eb3f7 100644
--- a/tests/helpers/image_generation_test_helper.py
+++ b/tests/helpers/image_generation_test_helper.py
@@ -24,36 +24,38 @@ class ImageGeneratorTestHelper:
         output_image_path = ImageGeneratorTestHelper.resolve_path(output_image_path)
         lora_paths = [str(ImageGeneratorTestHelper.resolve_path(p)) for p in lora_paths] if lora_paths else None
 
-        # given
-        flux = Flux1(
-            model_config=model_config,
-            quantize=8,
-            lora_paths=lora_paths,
-            lora_scales=lora_scales
-        )  # fmt: off
+        try:
+            # given
+            flux = Flux1(
+                model_config=model_config,
+                quantize=8,
+                lora_paths=lora_paths,
+                lora_scales=lora_scales
+            )  # fmt: off
 
-        # when
-        image = flux.generate_image(
-            seed=seed,
-            prompt=prompt,
-            config=Config(
-                num_inference_steps=steps,
-                height=341,
-                width=768,
-            ),
-        )
-        image.save(path=output_image_path)
+            # when
+            image = flux.generate_image(
+                seed=seed,
+                prompt=prompt,
+                config=Config(
+                    num_inference_steps=steps,
+                    height=341,
+                    width=768,
+                ),
+            )
+            image.save(path=output_image_path)
 
-        # then
-        np.testing.assert_array_equal(
-            np.array(Image.open(output_image_path)),
-            np.array(Image.open(reference_image_path)),
-            err_msg="Generated image doesn't match reference image",
-        )
+            # then
+            np.testing.assert_array_equal(
+                np.array(Image.open(output_image_path)),
+                np.array(Image.open(reference_image_path)),
+                err_msg="Generated image doesn't match reference image",
+            )
 
-        # cleanup
-        if os.path.exists(output_image_path):
-            os.remove(output_image_path)
+        finally:
+            # cleanup
+            if os.path.exists(output_image_path):
+                os.remove(output_image_path)
 
     @staticmethod
     def resolve_path(path) -> Path: