🧪 test_calc_job: fix daemon-restart flake

src/aiida/manage/tests/pytest_fixtures.py

@@ -727,17 +727,18 @@ def submit_and_await(started_daemon_client):
     def _factory(
         submittable: Process | ProcessBuilder | ProcessNode,
         state: plumpy.ProcessState = plumpy.ProcessState.FINISHED,
-        timeout: int = 20,
+        timeout: float = 20,
         **kwargs,
     ):
         """Submit a process and wait for it to achieve the given state.
 
         :param submittable: A process, a process builder or a process node. If it is a process or builder, it is
             submitted first before awaiting the desired state.
         :param state: The process state to wait for, by default it waits for the submittable to be ``FINISHED``.
-        :param timeout: The time to wait for the process to achieve the state.
+        :param timeout: The number of seconds to wait for the process to achieve the state.
         :param kwargs: If the ``submittable`` is a process class, it is instantiated with the ``kwargs`` as inputs.
-        :raises RuntimeError: If the process fails to achieve the specified state before the timeout expires.
+        :raises RuntimeError: If the process terminates in a state other than the one specified, or if it fails to
+            achieve the specified state before the timeout expires.
         """
         if inspect.isclass(submittable) and issubclass(submittable, Process):  # type: ignore[unreachable]
             node = submit(submittable, **kwargs)  # type: ignore[unreachable]
@@ -748,24 +749,35 @@ def _factory(
         else:
             raise ValueError(f'type of submittable `{type(submittable)}` is not supported.')
 
-        start_time = time.time()
+        start_time = time.monotonic()
+        # Terminal members of ``ProcessState``, mirroring ``ProcessNode.is_terminated``.
+        terminal_states = (plumpy.ProcessState.EXCEPTED, plumpy.ProcessState.FINISHED, plumpy.ProcessState.KILLED)
 
-        while node.process_state is not state:
-            if node.is_excepted:
-                raise RuntimeError(f'The process excepted: {node.exception}')
+        while True:
+            # Read ``process_state`` once per iteration: each access re-queries the database, so splitting the target
+            # and terminal checks across reads could straddle a transition and fail just as ``state`` is reached.
+            current_state = node.process_state
 
-            if time.time() - start_time >= timeout:
+            if current_state is state:
+                return node
+
+            # A terminal state cannot change, so waiting for a different one would only delay the failure.
+            if current_state in terminal_states:
+                if current_state is plumpy.ProcessState.EXCEPTED:
+                    raise RuntimeError(f'The process excepted: {node.exception}')
+                msg = f'The process terminated in state `{current_state}` while waiting for state `{state}`.'
+                raise RuntimeError(msg)
+
+            if time.monotonic() - start_time >= timeout:
                 daemon_log_file = pathlib.Path(started_daemon_client.daemon_log_file).read_text(encoding='utf-8')
                 daemon_status = 'running' if started_daemon_client.is_daemon_running else 'stopped'
                 raise RuntimeError(
-                    f'Timed out waiting for process with state `{node.process_state}` to enter state `{state}`.\n'
+                    f'Timed out waiting for process with state `{current_state}` to enter state `{state}`.\n'
                     f'Daemon <{started_daemon_client.profile.name}|{daemon_status}> log file content: \n'
                     f'{daemon_log_file}'
                 )
             time.sleep(0.1)
 
-        return node
-
     return _factory
 
 

src/aiida/tools/pytest_fixtures/daemon.py

@@ -113,17 +113,18 @@ def test(submit_and_await):
     :param submittable: A process, a process builder or a process node. If it is a process or builder, it is submitted
         first before awaiting the desired state.
     :param state: The process state to wait for, by default it waits for the submittable to be ``FINISHED``.
-    :param timeout: The time to wait for the process to achieve the state.
+    :param timeout: The number of seconds to wait for the process to achieve the state.
     :param kwargs: If the ``submittable`` is a process class, it is instantiated with the ``kwargs`` as inputs.
-    :raises RuntimeError: If the process fails to achieve the specified state before the timeout expires.
+    :raises RuntimeError: If the process terminates in a state other than the one specified, or if it fails to
+        achieve the specified state before the timeout expires.
     :returns `~aiida.orm.nodes.process.process.ProcessNode`: The process node.
     """
     from aiida.engine import ProcessState
 
     def factory(
         submittable: type[Process] | ProcessBuilder | ProcessNode,
         state: ProcessState = ProcessState.FINISHED,
-        timeout: int = 20,
+        timeout: float = 20,
         **kwargs,
     ):
         import inspect
@@ -141,22 +142,33 @@ def factory(
         else:
             raise ValueError(f'type of submittable `{type(submittable)}` is not supported.')
 
-        start_time = time.time()
+        start_time = time.monotonic()
+        # Terminal members of ``ProcessState``, mirroring ``ProcessNode.is_terminated``.
+        terminal_states = (ProcessState.EXCEPTED, ProcessState.FINISHED, ProcessState.KILLED)
 
-        while node.process_state is not state:
-            if node.is_excepted:
-                raise RuntimeError(f'The process excepted: {node.exception}')
+        while True:
+            # Read ``process_state`` once per iteration: each access re-queries the database, so splitting the target
+            # and terminal checks across reads could straddle a transition and fail just as ``state`` is reached.
+            current_state = node.process_state
 
-            if time.time() - start_time >= timeout:
+            if current_state is state:
+                return node
+
+            # A terminal state cannot change, so waiting for a different one would only delay the failure.
+            if current_state in terminal_states:
+                if current_state is ProcessState.EXCEPTED:
+                    raise RuntimeError(f'The process excepted: {node.exception}')
+                msg = f'The process terminated in state `{current_state}` while waiting for state `{state}`.'
+                raise RuntimeError(msg)
+
+            if time.monotonic() - start_time >= timeout:
                 daemon_log_file = pathlib.Path(started_daemon_client.daemon_log_file).read_text(encoding='utf-8')
                 daemon_status = 'running' if started_daemon_client.is_daemon_running else 'stopped'
                 raise RuntimeError(
-                    f'Timed out waiting for process with state `{node.process_state}` to enter state `{state}`.\n'
+                    f'Timed out waiting for process with state `{current_state}` to enter state `{state}`.\n'
                     f'Daemon <{started_daemon_client.profile.name}|{daemon_status}> log file content: \n'
                     f'{daemon_log_file}'
                 )
             time.sleep(0.1)
 
-        return node
-
     return factory

tests/engine/processes/calcjobs/test_calc_job.py

@@ -1348,8 +1348,6 @@ def test_restart_after_daemon_reset(get_calcjob_builder, daemon_client, submit_a
 
     This is a regression test for https://github.com/aiidateam/aiida-core/issues/5882.
     """
-    import time
-
     import plumpy
 
     daemon_client.start_daemon()
@@ -1363,19 +1361,10 @@ def test_restart_after_daemon_reset(get_calcjob_builder, daemon_client, submit_a
 
     daemon_client.restart_daemon(wait=True)
 
-    start_time = time.time()
-    timeout = 10
-
-    while node.process_state not in [plumpy.ProcessState.FINISHED, plumpy.ProcessState.EXCEPTED]:
-        if node.is_excepted:
-            raise AssertionError(f'The process excepted: {node.exception}')
-
-        if time.time() - start_time >= timeout:
-            raise AssertionError(f'process failed to terminate within timeout, current state: {node.process_state}')
-
-        time.sleep(0.1)
+    # The full stop/restart/reload/resume/finish cycle is heavy, so allow more than the fixture default: under CPU
+    # contention the post-restart wait alone approaches the 10 seconds this test used to allow.
+    submit_and_await(node, plumpy.ProcessState.FINISHED, timeout=30)
 
-    assert node.is_finished, node.process_state
     assert node.is_finished_ok, node.exit_status