From 92f5800451fb40b42f64b262abba6ffd3e339f14 Mon Sep 17 00:00:00 2001
From: Zoheb Shaikh <zoheb.shaikh@diamond.ac.uk>
Date: Fri, 17 Jan 2025 17:11:12 +0000
Subject: [PATCH] Made improvements

---
 src/blueapi/cli/cli.py                    |  2 ++
 src/blueapi/client/client.py              |  4 ++--
 src/blueapi/service/runner.py             | 12 ++++++------
 tests/system_tests/test_blueapi_system.py |  4 +++-
 tests/unit_tests/client/test_client.py    |  4 +++-
 5 files changed, 16 insertions(+), 10 deletions(-)

diff --git a/src/blueapi/cli/cli.py b/src/blueapi/cli/cli.py
index a678f7e7a..16046dd4d 100644
--- a/src/blueapi/cli/cli.py
+++ b/src/blueapi/cli/cli.py
@@ -343,8 +343,10 @@ def env(
         # Reload the environment if needed
         print("Reloading environment")
         status = client.reload_environment(timeout=timeout)
+        print("Environment is initialized")
     else:
         status = client.get_environment()
+
     print(status)
 
 
diff --git a/src/blueapi/client/client.py b/src/blueapi/client/client.py
index e732b4054..5faac378b 100644
--- a/src/blueapi/client/client.py
+++ b/src/blueapi/client/client.py
@@ -417,9 +417,9 @@ def _wait_for_reload(
                     f"Error reloading environment: {status.error_message}"
                 )
             elif (
-                status.initialized and status.environment_id != previous_environment_id
+                status.initialized is True
+                and status.environment_id != previous_environment_id
             ):
-                print("Environment is initialized")
                 return status
             time.sleep(polling_interval)
         # If the function did not raise or return early, it timed out
diff --git a/src/blueapi/service/runner.py b/src/blueapi/service/runner.py
index 508ee3677..734d24bf8 100644
--- a/src/blueapi/service/runner.py
+++ b/src/blueapi/service/runner.py
@@ -71,17 +71,17 @@ def reload(self):
 
     @start_as_current_span(TRACER)
     def start(self):
-        new_environment_id = uuid.uuid4()
+        environment_id = uuid.uuid4()
         try:
             self._subprocess = self._subprocess_factory()
             self.run(setup, self._config)
             self._state = EnvironmentResponse(
-                environment_id=new_environment_id,
+                environment_id=environment_id,
                 initialized=True,
             )
         except Exception as e:
             self._state = EnvironmentResponse(
-                environment_id=new_environment_id,
+                environment_id=environment_id,
                 initialized=False,
                 error_message=str(e),
             )
@@ -89,20 +89,20 @@ def start(self):
 
     @start_as_current_span(TRACER)
     def stop(self):
-        existing_environment_id = self._state.environment_id
+        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,
+                environment_id=environment_id,
                 initialized=False,
                 error_message=self._state.error_message,
             )
         except Exception as e:
             self._state = EnvironmentResponse(
-                environment_id=existing_environment_id,
+                environment_id=environment_id,
                 initialized=False,
                 error_message=str(e),
             )
diff --git a/tests/system_tests/test_blueapi_system.py b/tests/system_tests/test_blueapi_system.py
index 39beef408..c300a6169 100644
--- a/tests/system_tests/test_blueapi_system.py
+++ b/tests/system_tests/test_blueapi_system.py
@@ -338,5 +338,7 @@ def test_get_current_state_of_environment(client: BlueapiClient):
 
 
 def test_delete_current_environment(client: BlueapiClient):
+    current_env = client.get_environment()
     client.reload_environment()
-    assert client.get_environment().initialized
+    new_env = client.get_environment()
+    assert new_env.initialized and new_env != current_env
diff --git a/tests/unit_tests/client/test_client.py b/tests/unit_tests/client/test_client.py
index 27e36e915..99323b02d 100644
--- a/tests/unit_tests/client/test_client.py
+++ b/tests/unit_tests/client/test_client.py
@@ -78,7 +78,9 @@ def mock_rest() -> BlueapiRestClient:
     mock.get_all_tasks.return_value = TASKS
     mock.get_active_task.return_value = ACTIVE_TASK
     mock.get_environment.return_value = ENV
-    mock.delete_environment.return_value = ENV
+    mock.delete_environment.return_value = EnvironmentResponse(
+        environment_id=ENVIRONMENT_ID, initialized=False
+    )
     return mock