Using the Lightning CLI interface to run tests

Signed-off-by: João Lucas de Sousa Almeida <[email protected]>

Author: Joao-L-S-Almeida

tests/test_finetune.py

@@ -6,8 +6,7 @@
 import subprocess
 import os 
 
-from terratorch.models.backbones.prithvi_vit import checkpoint_filter_fn as checkpoint_filter_fn_vit
-from terratorch.models.backbones.prithvi_swin import checkpoint_filter_fn as checkpoint_filter_fn_swin
+from terratorch.cli_tools import build_lightning_cli
 
 @pytest.mark.parametrize("model_name", ["prithvi_swin_B", "prithvi_swin_L", "prithvi_vit_100", "prithvi_vit_300"])
 def test_finetune_multiple_backbones(model_name):
@@ -21,11 +20,7 @@ def test_finetune_multiple_backbones(model_name):
     torch.save(state_dict, os.path.join("tests/", model_name + ".pt"))
 
     # Running the terratorch CLI
-    command_str = f"terratorch fit -c tests/manufactured-finetune_{model_name}.yaml" 
-
-    command_out = subprocess.run(command_str, shell=True) 
-
-    assert not command_out.returncode
-
+    command_list = ["fit", "-c", f"tests/manufactured-finetune_{model_name}.yaml"]
+    _ = build_lightning_cli(command_list)