Block until environment reloaded

src/blueapi/service/model.py

@@ -1,5 +1,6 @@
 from collections.abc import Iterable
 from typing import Any
+from uuid import UUID
 
 from bluesky.protocols import HasName
 from pydantic import Field
@@ -144,6 +145,10 @@ class EnvironmentResponse(BlueapiBaseModel):
     State of internal environment.
     """
 
+    environment_id: UUID = Field(
+        description="Unique ID for the environment instance, can be used to "
+        "differentiate between a new environment and old that has been torn down"
+    )
     initialized: bool = Field(description="blueapi context initialized")
     error_message: str | None = Field(
         default=None,

src/blueapi/service/runner.py

@@ -1,6 +1,7 @@
 import inspect
 import logging
 import signal
+import uuid
 from collections.abc import Callable
 from importlib import import_module
 from multiprocessing import Pool, set_start_method
@@ -57,6 +58,7 @@ def default_subprocess_factory():
         self._subprocess = None
         self._subprocess_factory = subprocess_factory or default_subprocess_factory
         self._state = EnvironmentResponse(
+            environment_id=uuid.uuid4(),
             initialized=False,
         )
 
@@ -69,30 +71,38 @@ def reload(self):
 
     @start_as_current_span(TRACER)
     def start(self):
+        new_environment_id = uuid.uuid4()
         try:
             self._subprocess = self._subprocess_factory()
             self.run(setup, self._config)
-            self._state = EnvironmentResponse(initialized=True)
+            self._state = EnvironmentResponse(
+                environment_id=new_environment_id,
+                initialized=True,
+            )
         except Exception as e:
             self._state = EnvironmentResponse(
+                environment_id=new_environment_id,
                 initialized=False,
                 error_message=str(e),
             )
             LOGGER.exception(e)
 
     @start_as_current_span(TRACER)
     def stop(self):
+        existing_environment_id = self._state.environment_id
         try:
             self.run(teardown)
             if self._subprocess is not None:
                 self._subprocess.close()
                 self._subprocess.join()
             self._state = EnvironmentResponse(
+                environment_id=existing_environment_id,
                 initialized=False,
                 error_message=self._state.error_message,
             )
         except Exception as e:
             self._state = EnvironmentResponse(
+                environment_id=existing_environment_id,
                 initialized=False,
                 error_message=str(e),
             )

tests/unit_tests/client/test_client.py

@@ -1,5 +1,5 @@
 from collections.abc import Callable
-from unittest.mock import MagicMock, Mock, call
+from unittest.mock import MagicMock, Mock, call, patch
 
 import pytest
 from bluesky_stomp.messaging import MessageContext
@@ -259,6 +259,72 @@ def test_reload_environment(
     mock_rest.delete_environment.assert_called_once()
 
 
+@patch("blueapi.client.client.time.time")
+@patch("blueapi.client.client.time.sleep")
+def test_reload_environment_no_timeout(
+    mock_sleep: Mock,
+    mock_time: Mock,
+    client: BlueapiClient,
+    mock_rest: Mock,
+):
+    mock_rest.get_environment.side_effect = [
+        EnvironmentResponse(initialized=False),
+        EnvironmentResponse(initialized=False),
+        EnvironmentResponse(initialized=False),
+        EnvironmentResponse(initialized=True),
+    ]
+    mock_time.return_value = 100.0
+    client.reload_environment(timeout=None)
+    assert mock_sleep.call_count == 3
+
+
+@patch("blueapi.client.client.time.time")
+@patch("blueapi.client.client.time.sleep")
+def test_reload_environment_with_timeout(
+    _: Mock,
+    mock_time: Mock,
+    client: BlueapiClient,
+    mock_rest: Mock,
+):
+    mock_rest.get_environment.side_effect = [
+        EnvironmentResponse(initialized=False),
+        EnvironmentResponse(initialized=False),
+        EnvironmentResponse(initialized=False),
+        EnvironmentResponse(initialized=False),
+    ]
+    mock_time.side_effect = [
+        100.0,
+        100.5,
+        101.0,  # Timeout should occur here
+        101.5,
+    ]
+    with pytest.raises(
+        TimeoutError,
+        match="Failed to reload the environment within 1.0 "
+        "seconds, a server restart is recommended",
+    ):
+        client.reload_environment(timeout=1.0)
+
+
+@patch("blueapi.client.client.time.time")
+@patch("blueapi.client.client.time.sleep")
+def test_reload_environment_ignores_current_environment(
+    mock_sleep: Mock,
+    mock_time: Mock,
+    client: BlueapiClient,
+    mock_rest: Mock,
+):
+    mock_rest.get_environment.side_effect = [
+        EnvironmentResponse(initialized=True),  # This is the old environment
+        EnvironmentResponse(initialized=False),
+        EnvironmentResponse(initialized=False),
+        EnvironmentResponse(initialized=True),  # This is the new environment
+    ]
+    mock_time.return_value = 100.0
+    client.reload_environment(timeout=None)
+    assert mock_sleep.call_count == 3
+
+
 def test_reload_environment_failure(
     client: BlueapiClient,
     mock_rest: Mock,