fix finetune tests

Signed-off-by: Carlos Gomes <[email protected]>

tests/test_finetune.py

@@ -1,6 +1,5 @@
-import importlib
 import os
-import subprocess
+import shutil
 
 import pytest
 import timm
@@ -10,47 +9,31 @@
 
 
 @pytest.fixture(autouse=True)
-def cleanup():
-    yield # everything after this runs after each test
-
-    for file in os.listdir("tests"):
-        if file.endswith(".pt"):
-            os.remove(os.path.join("tests", file))
-
-@pytest.mark.parametrize("model_name", ["prithvi_swin_B", "prithvi_swin_L", "prithvi_vit_100", "prithvi_vit_300"])
-def test_finetune_multiple_backbones(model_name):
+def setup_and_cleanup(model_name):
     model_instance = timm.create_model(model_name)
 
     state_dict = model_instance.state_dict()
 
     torch.save(state_dict, os.path.join("tests", model_name + ".pt"))
 
-    # Running the terratorch CLI
+    yield # everything after this runs after each test
+
+    os.remove(os.path.join("tests", model_name + ".pt"))
+    shutil.rmtree(os.path.join("tests", "all_ecos_random"))
+
+@pytest.mark.parametrize("model_name", ["prithvi_swin_B", "prithvi_swin_L", "prithvi_vit_100", "prithvi_vit_300"])
+def test_finetune_multiple_backbones(model_name):
     command_list = ["fit", "-c", f"tests/manufactured-finetune_{model_name}.yaml"]
     _ = build_lightning_cli(command_list)
 
 
 @pytest.mark.parametrize("model_name", ["prithvi_swin_B"])
 def test_finetune_bands_intervals(model_name):
-    model_instance = timm.create_model(model_name)
-
-    state_dict = model_instance.state_dict()
-
-    torch.save(state_dict, os.path.join("tests", model_name + ".pt"))
-
-    # Running the terratorch CLI
     command_list = ["fit", "-c", f"tests/manufactured-finetune_{model_name}_band_interval.yaml"]
     _ = build_lightning_cli(command_list)
 
 
 @pytest.mark.parametrize("model_name", ["prithvi_swin_B"])
 def test_finetune_bands_str(model_name):
-    model_instance = timm.create_model(model_name)
-
-    state_dict = model_instance.state_dict()
-
-    torch.save(state_dict, os.path.join("tests", model_name + ".pt"))
-
-    # Running the terratorch CLI
     command_list = ["fit", "-c", f"tests/manufactured-finetune_{model_name}_string.yaml"]
     _ = build_lightning_cli(command_list)