Apply review comments

Author: addyess

src/reactive/containerd.py

@@ -532,9 +532,13 @@ def _apt_restart_services(restart: bool):
     if not restart:
         log("Services will be not restarted after apt operations.")
         os.environ.update({"NEEDRESTART_SUSPEND": "1"})
-    yield
-    if restore is not None:
-        os.environ["NEEDRESTART_SUSPEND"] = restore
+    try:
+        yield
+    finally:
+        if restore is not None:
+            os.environ["NEEDRESTART_SUSPEND"] = restore
+        else:
+            os.environ.pop("NEEDRESTART_SUSPEND", None)
 
 
 def reinstall_containerd(restart: bool = True):

tests/unit/test_containerd_reactive.py

@@ -410,12 +410,12 @@ def test_needs_gpu_reboot_true(check_output, is_state, remove_state, set_state):
 
 
 @pytest.mark.parametrize(
-    "restart,initial_env,expected_during,expected_after,log_called",
+    "restart,initial_env,expected_during,expected_after",
     [
-        (True, {}, None, None, False),
-        (False, {}, "1", "1", True),
-        (True, {"NEEDRESTART_SUSPEND": "original_value"}, None, "original_value", False),
-        (False, {"NEEDRESTART_SUSPEND": "original_value"}, "1", "original_value", True),
+        (True, {}, None, None),
+        (False, {}, "1", None),
+        (True, {"NEEDRESTART_SUSPEND": "original"}, None, "original"),
+        (False, {"NEEDRESTART_SUSPEND": "original"}, "1", "original"),
     ],
     ids=[
         "restart=True, no initial env",
@@ -426,9 +426,10 @@ def test_needs_gpu_reboot_true(check_output, is_state, remove_state, set_state):
 )
 @mock.patch.object(containerd, "log")
 @mock.patch.dict(os.environ, {}, clear=True)
-def test_apt_restart_services(mock_log, restart, initial_env, expected_during, expected_after, log_called):
+def test_apt_restart_services(mock_log, restart, initial_env, expected_during, expected_after):
     """Verify _apt_restart_services behavior with various configurations."""
     # Setup initial environment from parameters
+    log_called = not restart
     os.environ.update(initial_env)
 
     with containerd._apt_restart_services(restart=restart):