From 2f6b1a3eb89d4cf555894a65e78b0589658a89d4 Mon Sep 17 00:00:00 2001
From: yuvrajjain2003 <yjain@unc.edu>
Date: Fri, 1 Dec 2023 00:28:47 -0500
Subject: [PATCH] Implement the new Job Manager in ROI Extractor

---
 clpipe/job_manager.py       |  2 +-
 clpipe/roi_extractor.py     | 24 +++++++++++++-----------
 tests/test_batch_manager.py |  4 ++--
 3 files changed, 16 insertions(+), 14 deletions(-)

diff --git a/clpipe/job_manager.py b/clpipe/job_manager.py
index b5392475..aa4d9dc3 100644
--- a/clpipe/job_manager.py
+++ b/clpipe/job_manager.py
@@ -248,7 +248,7 @@ def add_job(self, job):
     def compilejobstrings(self):
         header = self.createsubmissionhead()
         for job in self.jobs:
-            temp = header.format(jobid=job.jobID, cmdwrap=job.jobString)
+            temp = header.format(jobid=job.job_name, cmdwrap=job.job_string)
             self.submission_list.append(temp)
 
     def compile_job_strings(self):
diff --git a/clpipe/roi_extractor.py b/clpipe/roi_extractor.py
index d27a2750..130b801f 100644
--- a/clpipe/roi_extractor.py
+++ b/clpipe/roi_extractor.py
@@ -14,7 +14,7 @@
 import glob
 import shutil
 from .config.options import ProjectOptions
-from .job_manager import BatchManager, Job
+from .job_manager import JobManagerFactory
 from pkg_resources import resource_stream, resource_filename
 from .errors import MaskFileNotFoundError
 from .utils import get_logger, resolve_fmriprep_dir
@@ -90,12 +90,15 @@ def fmri_roi_extraction(
         "-custom_type={custom_type} -single"
     )
 
-    batch_manager = BatchManager(config.batch_config_path, config.roi_extraction.log_directory)
-    batch_manager.update_mem_usage(config.roi_extraction.memory_usage)
-    batch_manager.update_time(config.roi_extraction.time_usage)
-    batch_manager.update_nthreads(config.roi_extraction.n_threads)
-    batch_manager.update_email(config.email_address)
-    batch_manager.createsubmissionhead()
+    batch_manager = JobManagerFactory.get(
+        batch_config=config.batch_config_path,
+        output_directory=config.roi_extraction.log_directory,
+        mem_use=config.roi_extraction.time_usage,
+        time=config.roi_extraction.time_usage,
+        threads=config.roi_extraction.n_threads,
+        email=config.email_address
+    )
+    
     for subject in sublist:
         logger.debug(f"Setting up ROI extraction for subject {subject}")
         for cur_atlas in atlas_list:
@@ -172,8 +175,9 @@ def fmri_roi_extraction(
                 logger.debug("Overlap ok flag set")
 
             sub_string_temp = sub_string_temp + " " + subject
-            batch_manager.addjob(
-                Job("ROI_extract_" + subject + "_" + atlas_name, sub_string_temp)
+            batch_manager.add_job(
+                "ROI_extract_" + subject + "_" + atlas_name, 
+                sub_string_temp
             )
             if single:
                 _fmri_roi_extract_subject(
@@ -192,10 +196,8 @@ def fmri_roi_extraction(
                 )
     if not single:
         if submit:
-            batch_manager.compilejobstrings()
             batch_manager.submit_jobs()
         else:
-            batch_manager.compilejobstrings()
             click.echo(batch_manager.print_jobs())
 
 
diff --git a/tests/test_batch_manager.py b/tests/test_batch_manager.py
index 1a9dc30a..59888998 100644
--- a/tests/test_batch_manager.py
+++ b/tests/test_batch_manager.py
@@ -15,14 +15,14 @@ def test_adding_jobs():
     batch_manager.add_job(job1)
 
     assert len(batch_manager.jobs) == 1
-    assert batch_manager.jobs[0].job_id == 1
+    assert batch_manager.jobs[0].job_name == 1
     assert batch_manager.jobs[0].job_string == "echo hi"
 
     job2 = Job(2, "additional_job")
     batch_manager.add_job(job2)
 
     assert len(batch_manager.jobs) == 2
-    assert batch_manager.jobs[1].job_id == 2
+    assert batch_manager.jobs[1].job_name == 2
     assert batch_manager.jobs[1].job_string == "additional_job"