Fix simulator tests by using the correct queue API

equinor · Nov 21, 2023 · e0029a7 · e0029a7
1 parent 83db7f0
commit e0029a7
Show file tree

Hide file tree

Showing 4 changed files with 30 additions and 26 deletions.
diff --git a/src/ert/simulator/batch_simulator_context.py b/src/ert/simulator/batch_simulator_context.py
@@ -55,11 +55,11 @@ def status(self) -> Status:
         Will return the state of the simulations.
         """
         return Status(
-            running=self._job_queue.count_status(RealizationState.RUNNING),
-            waiting=self._job_queue.count_status(RealizationState.WAITING),
-            pending=self._job_queue.count_status(RealizationState.PENDING),
-            complete=self._job_queue.count_status(RealizationState.SUCCESS),
-            failed=self._job_queue.count_status(RealizationState.FAILED),
+            running=self.job_queue.count_realization_state(RealizationState.RUNNING),
+            waiting=self.job_queue.count_realization_state(RealizationState.WAITING),
+            pending=self.job_queue.count_realization_state(RealizationState.PENDING),
+            complete=self.job_queue.count_realization_state(RealizationState.SUCCESS),
+            failed=self.job_queue.count_realization_state(RealizationState.FAILED),
         )
 
     def results(self) -> List[Optional[Dict[str, "npt.NDArray[np.float64]"]]]:

diff --git a/src/ert/simulator/simulation_context.py b/src/ert/simulator/simulation_context.py
@@ -54,7 +54,7 @@ def _run_forward_model(
     for index, run_arg in enumerate(run_context):
         if not run_context.is_active(index):
             continue
-        job_queue.add_job_from_run_arg(
+        job_queue.add_realization_from_run_arg(
             run_arg,
             ert.ert_config.queue_config.job_script,
             max_runtime,
@@ -182,10 +182,10 @@ def job_queue(self) -> JobQueue:
 
     def __repr__(self) -> str:
         running = "running" if self.isRunning() else "not running"
-        numRunn = self.job_queue.count_status(RealizationState.RUNNING)
-        numSucc = self.job_queue.count_status(RealizationState.SUCCESS)
-        numFail = self.job_queue.count_status(RealizationState.FAILED)
-        numWait = self.job_queue.count_status(RealizationState.WAITING)
+        numRunn = self.job_queue.count_realization_state(RealizationState.RUNNING)
+        numSucc = self.job_queue.count_realization_state(RealizationState.SUCCESS)
+        numFail = self.job_queue.count_realization_state(RealizationState.FAILED)
+        numWait = self.job_queue.count_realization_state(RealizationState.WAITING)
         return (
             f"SimulationContext({running}, #running = {numRunn}, "
             f"#success = {numSucc}, #failed = {numFail}, #waiting = {numWait})"

diff --git a/tests/unit_tests/simulator/test_batch_sim.py b/tests/unit_tests/simulator/test_batch_sim.py
@@ -5,7 +5,7 @@
 import pytest
 
 from ert.config import ErtConfig
-from ert.job_queue import JobStatus
+from ert.realization_state import RealizationState
 from ert.simulator import BatchSimulator
 
 
@@ -409,13 +409,13 @@ def test_stop_sim(copy_case, storage):
 def assertContextStatusOddFailures(batch_ctx, final_state_only=False):
     running_status = set(
         (
-            JobStatus.WAITING,
-            JobStatus.SUBMITTED,
-            JobStatus.PENDING,
-            JobStatus.RUNNING,
-            JobStatus.UNKNOWN,
-            JobStatus.EXIT,
-            JobStatus.DONE,
+            RealizationState.WAITING,
+            RealizationState.SUBMITTED,
+            RealizationState.PENDING,
+            RealizationState.RUNNING,
+            RealizationState.UNKNOWN,
+            RealizationState.EXIT,
+            RealizationState.DONE,
             None,  # job is not submitted yet but ok for this test
         )
     )
@@ -425,9 +425,9 @@ def assertContextStatusOddFailures(batch_ctx, final_state_only=False):
         if not final_state_only and status in running_status:
             continue
         if idx % 2 == 0:
-            assert status == JobStatus.SUCCESS
+            assert status == RealizationState.SUCCESS
         else:
-            assert status == JobStatus.FAILED
+            assert status == RealizationState.FAILED
 
 
 def test_batch_ctx_status_failing_jobs(setup_case, storage):

diff --git a/tests/unit_tests/simulator/test_simulation_context.py b/tests/unit_tests/simulator/test_simulation_context.py
@@ -46,13 +46,17 @@ def test_simulation_context(setup_case, storage):
                 f"runpath/realization-{iens}-{iens}/iter-0"
             )
 
-    assert even_ctx.job_queue.count_status(RealizationState.FAILED) == 0
-    assert even_ctx.job_queue.count_status(RealizationState.RUNNING) == 0
-    assert even_ctx.job_queue.count_status(RealizationState.SUCCESS) == size / 2
+    assert even_ctx.job_queue.count_realization_state(RealizationState.FAILED) == 0
+    assert even_ctx.job_queue.count_realization_state(RealizationState.RUNNING) == 0
+    assert (
+        even_ctx.job_queue.count_realization_state(RealizationState.SUCCESS) == size / 2
+    )
 
-    assert odd_ctx.job_queue.count_status(RealizationState.FAILED) == 0
-    assert odd_ctx.job_queue.count_status(RealizationState.RUNNING) == 0
-    assert odd_ctx.job_queue.count_status(RealizationState.SUCCESS) == size / 2
+    assert odd_ctx.job_queue.count_realization_state(RealizationState.FAILED) == 0
+    assert odd_ctx.job_queue.count_realization_state(RealizationState.RUNNING) == 0
+    assert (
+        odd_ctx.job_queue.count_realization_state(RealizationState.SUCCESS) == size / 2
+    )
 
     for iens in range(size):
         if iens % 2 == 0: