additional tests, plus coverage adjusts, plus remove an exception block that can't happen

Author: jlnav

libensemble/executors/flux_executor.py

@@ -166,27 +166,19 @@ def wait(self, timeout: float | None = None) -> None:
         if not self._check_poll():
             return
 
-        try:
-            # Wait for job to complete
-            start_time = time.time()
-            while True:
-                self.poll()
-                if self.finished:
-                    break
-
-                if timeout is not None:
-                    elapsed = time.time() - start_time
-                    if elapsed >= timeout:
-                        raise TimeoutExpired(self.name, timeout)
-
-                time.sleep(0.1)
-
-        except TimeoutExpired:
-            raise
-        except Exception as e:
-            logger.warning(f"Error waiting for Flux job {self.flux_jobid}: {e}")
-            self.state = "FAILED"
-            self.finished = True
+        # Wait for job to complete
+        start_time = time.time()
+        while True:
+            self.poll()
+            if self.finished:
+                break
+
+            if timeout is not None:
+                elapsed = time.time() - start_time
+                if elapsed >= timeout:
+                    raise TimeoutExpired(self.name, timeout)
+
+            time.sleep(0.1)
 
     def kill(self, wait_time: int | None = 60) -> None:
         """Kills/cancels the Flux job.
@@ -205,10 +197,6 @@ def kill(self, wait_time: int | None = 60) -> None:
             logger.warning(f"Trying to kill task that is no longer running. Task {self.name}: Status is {self.state}")
             return
 
-        if self.flux_jobid is None:
-            logger.warning(f"Task {self.name} has no Flux job ID - cannot kill")
-            return
-
         logger.info(f"Canceling Flux job {self.flux_jobid} for task {self.name}")
 
         try:
@@ -287,14 +275,6 @@ def __init__(self) -> None:
         self.resources = None
         self.platform_info: dict = {}
 
-    def set_resources(self, resources) -> None:
-        """Set resources for the executor."""
-        self.resources = resources
-
-    def add_platform_info(self, platform_info: dict | None = None) -> None:
-        """Add platform info to the executor."""
-        self.platform_info = platform_info or {}
-
     def submit(
         self,
         calc_type: str | None = None,

libensemble/tests/unit_tests/test_flux.py

@@ -398,7 +398,6 @@ def test_flux_executor_submit_builds_jobspec_with_environment_and_gpus():
 
     executor = object.__new__(flux_executor.FluxExecutor)
     executor.flux_handle = object()
-    executor.resources = None
     executor.platform_info = {}
     executor.workerID = 7
     executor.list_of_tasks = []
@@ -463,7 +462,6 @@ def test_flux_executor_init_connects_with_flux_uri():
 
     fake_flux_module.Flux.assert_called_once_with()
     assert executor.flux_handle == "flux-handle"
-    assert executor.resources is None
     assert executor.platform_info == {}
 
 
@@ -597,6 +595,52 @@ def test_flux_task_set_complete_handles_dry_run_and_return_codes():
     assert finished_task.success is False
     assert finished_task.state == "FAILED"
 
+    # cover waiting on a task that completes before timeout
+    task = flux_executor.FluxTask(
+        app=SimpleNamespace(name="app"),
+        app_args=None,
+        workdir=os.getcwd(),
+        stdout="out.txt",
+        stderr="err.txt",
+        workerid=1,
+        dry_run=False,
+    )
+    task._set_complete()
+    task.flux_jobid = 123
+    task.wait(timeout=10)
+    task.kill()
+
+
+def test_flux_task_dry_run_exception_and_kill():
+    """Test FluxTask dry run exception attributes."""
+    task = flux_executor.FluxTask(
+        app=SimpleNamespace(name="app"),
+        app_args=None,
+        workdir=os.getcwd(),
+        stdout="out.txt",
+        stderr="err.txt",
+        workerid=1,
+        dry_run=True,
+    )
+    task.wait()
+    assert task.finished is True
+    assert task.success is True
+    assert task.state == "FINISHED"
+    task.kill()
+    task = flux_executor.FluxTask(
+        app=SimpleNamespace(name="app"),
+        app_args=None,
+        workdir=os.getcwd(),
+        stdout="out.txt",
+        stderr="err.txt",
+        workerid=1,
+        dry_run=True,
+    )
+    task.poll()
+    assert task.finished is True
+    assert task.success is True
+    assert task.state == "FINISHED"
+
 
 def test_flux_task_wait_completes_and_times_out():
     """Test FluxTask wait completes after polling and raises on timeout."""
@@ -716,6 +760,7 @@ def poll_side_effect():
     # Validator tests
     test_validator_accepts_flux()
     test_validator_accepts_all_runners()
+    test_validator_rejects_invalid()
 
     # Platform tests
     test_flux_allocation_platform()
@@ -724,4 +769,14 @@ def poll_side_effect():
     # EnvResources tests
     test_env_resources_flux_env_variable()
 
+    # Flux Executor tests
+    test_flux_executor_init_connects_with_flux_uri()
+    test_flux_executor_wait_on_start_polls_until_running()
+    test_flux_task_poll_maps_completion_waiting_and_unknown_states()
+    test_flux_task_handle_completion_success_and_failure()
+    test_flux_task_set_complete_handles_dry_run_and_return_codes()
+    test_flux_task_dry_run_exception_and_kill()
+    test_flux_task_wait_completes_and_times_out()
+    test_flux_task_kill_cancels_and_marks_user_killed()
+
     print("All standalone tests passed!")