From a415baae0ffd9b3812f2cec7e799a38374a56685 Mon Sep 17 00:00:00 2001
From: Lin Guo <linsword13@gmail.com>
Date: Wed, 11 Dec 2024 15:08:10 +0000
Subject: [PATCH] Improve on test

---
 .../slurm_workflow_manager.py                 | 28 +++++++++++++++----
 1 file changed, 23 insertions(+), 5 deletions(-)

diff --git a/lib/ramble/ramble/test/workflow_manager_functionality/slurm_workflow_manager.py b/lib/ramble/ramble/test/workflow_manager_functionality/slurm_workflow_manager.py
index 30af0d334..a9e487f77 100644
--- a/lib/ramble/ramble/test/workflow_manager_functionality/slurm_workflow_manager.py
+++ b/lib/ramble/ramble/test/workflow_manager_functionality/slurm_workflow_manager.py
@@ -27,18 +27,19 @@ def test_slurm_workflow():
     test_config = """
 ramble:
   variants:
-    workflow_manager: slurm
+    workflow_manager: '{wm_name}'
   variables:
-    # This batch_submit is overridden
-    batch_submit: echo cmd_not_exist
+    # This batch_submit is overridden with slurm workflow manager
+    batch_submit: echo {wm_name}
     mpi_command: mpirun -n {n_ranks} -hostfile hostfile
     processes_per_node: 1
+    wm_name: ['None', 'slurm']
   applications:
     hostname:
       workloads:
         local:
           experiments:
-            test:
+            test_{wm_name}:
               variables:
                 n_nodes: 1
                 extra_sbatch_headers: "#SBATCH --gpus-per-task={n_threads}"
@@ -51,7 +52,24 @@ def test_slurm_workflow():
         ws._re_read()
         workspace("setup", "--dry-run", global_args=["-D", ws.root])
 
-        path = os.path.join(ws.experiment_dir, "hostname", "local", "test")
+        # assert the batch_submit is overriden, pointing to the generated script
+        all_exec_file = os.path.join(ws.root, "all_experiments")
+        with open(all_exec_file) as f:
+            content = f.read()
+            assert "echo None" in content
+            assert "echo slurm" not in content
+            assert os.path.join("hostname", "local", "test_slurm", "batch_submit") in content
+
+        # Assert on no workflow manager
+        path = os.path.join(ws.experiment_dir, "hostname", "local", "test_None")
+        files = [f for f in os.listdir(path) if os.path.isfile(os.path.join(path, f))]
+        assert "slurm_execute_experiment" not in files
+        assert "batch_submit" not in files
+        assert "query_job" not in files
+        assert "cancel_job" not in files
+
+        # Assert on slurm workflow manager
+        path = os.path.join(ws.experiment_dir, "hostname", "local", "test_slurm")
         files = [f for f in os.listdir(path) if os.path.isfile(os.path.join(path, f))]
         assert "slurm_execute_experiment" in files
         assert "batch_submit" in files