Make sure cleanup happens even if tests fail

This commit is contained in:
filipstrand 2024-10-09 23:38:56 +02:00
parent a9461cacf0
commit 78e2b726fb
2 changed files with 62 additions and 58 deletions

View File

@ -27,6 +27,7 @@ class ImageGeneratorControlnetTestHelper:
controlnet_image_path = str(ImageGeneratorTestHelper.resolve_path(controlnet_image_path)) 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 lora_paths = [str(ImageGeneratorTestHelper.resolve_path(p)) for p in lora_paths] if lora_paths else None
try:
# given # given
flux = Flux1Controlnet( flux = Flux1Controlnet(
model_config=model_config, model_config=model_config,
@ -58,6 +59,7 @@ class ImageGeneratorControlnetTestHelper:
err_msg="Generated image doesn't match reference image", err_msg="Generated image doesn't match reference image",
) )
finally:
# cleanup # cleanup
if os.path.exists(output_image_path): if os.path.exists(output_image_path):
os.remove(output_image_path) os.remove(output_image_path)

View File

@ -24,6 +24,7 @@ class ImageGeneratorTestHelper:
output_image_path = ImageGeneratorTestHelper.resolve_path(output_image_path) 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 lora_paths = [str(ImageGeneratorTestHelper.resolve_path(p)) for p in lora_paths] if lora_paths else None
try:
# given # given
flux = Flux1( flux = Flux1(
model_config=model_config, model_config=model_config,
@ -51,6 +52,7 @@ class ImageGeneratorTestHelper:
err_msg="Generated image doesn't match reference image", err_msg="Generated image doesn't match reference image",
) )
finally:
# cleanup # cleanup
if os.path.exists(output_image_path): if os.path.exists(output_image_path):
os.remove(output_image_path) os.remove(output_image_path)