From 06f336c71465928798100cf155494e7e40e62ded Mon Sep 17 00:00:00 2001
From: filipstrand <strand.filip@gmail.com>
Date: Sat, 1 Mar 2025 22:10:10 +0100
Subject: [PATCH] Pre test cleanup

---
 tests/dreambooth/test_resume_training.py        | 11 +++++++++++
 tests/dreambooth/test_train_and_load_weights.py | 12 ++++++++++++
 tests/model_saving/test_model_saving.py         | 12 ++++++++++++
 tests/model_saving/test_model_saving_lora.py    | 12 ++++++++++++
 4 files changed, 47 insertions(+)

diff --git a/tests/dreambooth/test_resume_training.py b/tests/dreambooth/test_resume_training.py
index 423119b..1d3c4c8 100644
--- a/tests/dreambooth/test_resume_training.py
+++ b/tests/dreambooth/test_resume_training.py
@@ -14,6 +14,9 @@ CHECKPOINT_5 = "tests/dreambooth/tmp/_checkpoints/0000005_checkpoint.zip"
 
 class TestResumeTraining:
     def test_resume_training(self):
+        # Clean up any existing temporary directories from previous test runs
+        TestResumeTraining.delete_folder_if_exists("tests/dreambooth/tmp")
+
         try:
             # Given: A small training run from scratch for 5 steps (as described in the config)...
             fluxA, runtime_config, training_spec, training_state = DreamBoothInitializer.initialize(
@@ -85,3 +88,11 @@ class TestResumeTraining:
     @staticmethod
     def delete_folder(path: str) -> None:
         return shutil.rmtree(path)
+
+    @staticmethod
+    def delete_folder_if_exists(path: str) -> None:
+        if os.path.exists(path):
+            shutil.rmtree(path)
+            print(f"Deleted folder: {path}")
+        else:
+            print("The specified folder does not exist.")
diff --git a/tests/dreambooth/test_train_and_load_weights.py b/tests/dreambooth/test_train_and_load_weights.py
index 58163e9..bb6456e 100644
--- a/tests/dreambooth/test_train_and_load_weights.py
+++ b/tests/dreambooth/test_train_and_load_weights.py
@@ -1,3 +1,4 @@
+import os
 import shutil
 
 import numpy as np
@@ -14,6 +15,9 @@ LORA_FILE = "tests/dreambooth/tmp/_checkpoints/0000005_checkpoint/0000005_adapte
 
 class TestTrainAndLoadWeights:
     def test_train_and_load_weights(self):
+        # Clean up any existing temporary directories from previous test runs
+        TestTrainAndLoadWeights.delete_folder_if_exists("tests/dreambooth/tmp")
+
         try:
             # Given: A small training run from scratch for 5 steps (as described in the config)...
             fluxA, runtime_config, training_spec, training_state = DreamBoothInitializer.initialize(
@@ -72,3 +76,11 @@ class TestTrainAndLoadWeights:
     @staticmethod
     def delete_folder(path: str) -> None:
         return shutil.rmtree(path)
+
+    @staticmethod
+    def delete_folder_if_exists(path: str) -> None:
+        if os.path.exists(path):
+            shutil.rmtree(path)
+            print(f"Deleted folder: {path}")
+        else:
+            print("The specified folder does not exist.")
diff --git a/tests/model_saving/test_model_saving.py b/tests/model_saving/test_model_saving.py
index 82a0e7a..96cc696 100644
--- a/tests/model_saving/test_model_saving.py
+++ b/tests/model_saving/test_model_saving.py
@@ -1,3 +1,4 @@
+import os
 import shutil
 
 import numpy as np
@@ -9,6 +10,9 @@ PATH = "tests/4bit/"
 
 class TestModelSaving:
     def test_save_and_load_4bit_model(self):
+        # Clean up any existing temporary directories from previous test runs
+        TestModelSaving.delete_folder_if_exists(PATH)
+
         try:
             # given a saved quantized model (and an image from that model)
             fluxA = Flux1(
@@ -56,3 +60,11 @@ class TestModelSaving:
     @staticmethod
     def delete_folder(path: str) -> None:
         return shutil.rmtree(path)
+
+    @staticmethod
+    def delete_folder_if_exists(path: str) -> None:
+        if os.path.exists(path):
+            shutil.rmtree(path)
+            print(f"Deleted folder: {path}")
+        else:
+            print(f"Folder does not exist: {path}")
diff --git a/tests/model_saving/test_model_saving_lora.py b/tests/model_saving/test_model_saving_lora.py
index acf3a88..88bf8ab 100644
--- a/tests/model_saving/test_model_saving_lora.py
+++ b/tests/model_saving/test_model_saving_lora.py
@@ -1,3 +1,4 @@
+import os
 import shutil
 from pathlib import Path
 
@@ -10,6 +11,9 @@ PATH = "tests/4bit/"
 
 class TestModelSavingLora:
     def test_save_and_load_4bit_model_with_lora(self):
+        # Clean up any existing temporary directories from previous test runs
+        TestModelSavingLora.delete_folder_if_exists(PATH)
+
         try:
             # given a saved quantized model on disk (without LoRA)...
             fluxA = Flux1(
@@ -71,6 +75,14 @@ class TestModelSavingLora:
     def delete_folder(path: str) -> None:
         return shutil.rmtree(path)
 
+    @staticmethod
+    def delete_folder_if_exists(path: str) -> None:
+        if os.path.exists(path):
+            shutil.rmtree(path)
+            print(f"Deleted folder: {path}")
+        else:
+            print(f"Folder does not exist: {path}")
+
     @staticmethod
     def resolve_path(path) -> Path | None:
         if path is None: