Make env var a literal value

Author: addyess

src/reactive/containerd.py

@@ -52,6 +52,7 @@
 from charmhelpers.fetch.ubuntu_apt_pkg import Package
 
 NVIDIA_SOURCES_FILE = "/etc/apt/sources.list.d/nvidia.list"
+NEEDRESTART_SUSPEND = "NEEDRESTART_SUSPEND"
 
 
 def apt_packages(packages: typing.Set[str]) -> typing.Mapping[str, Package]:
@@ -526,22 +527,21 @@ def _apt_restart_services(restart: bool):
     """
     Context manager to conditionally restart services after apt operations.
 
-    Args:
-        restart: whether to restart services after apt operations
+    :param bool restart: whether to restart services after apt operations
     """
-    original = os.environ.get("NEEDRESTART_SUSPEND")
-    log("Services will {}be restarted after apt operations.".format("" if restart else "not "))
+    env = NEEDRESTART_SUSPEND
+    original = os.environ.pop(env, None)
+    log(f"Services will {'' if restart else 'not '}be restarted after apt operations.")
     if not restart:
-        os.environ.update({"NEEDRESTART_SUSPEND": "1"})
-    else:
-        os.environ.pop("NEEDRESTART_SUSPEND", None)
+        os.environ.update({env: "1"})
+
     try:
         yield
     finally:
         if original is None:
-            os.environ.pop("NEEDRESTART_SUSPEND", None)
+            os.environ.pop(env, None)
         else:
-            os.environ["NEEDRESTART_SUSPEND"] = original
+            os.environ[env] = original
 
 
 def reinstall_containerd(restart: bool = True):

tests/unit/test_containerd_reactive.py

@@ -414,8 +414,8 @@ def test_needs_gpu_reboot_true(check_output, is_state, remove_state, set_state):
     [
         (True, {}, None, None),
         (False, {}, "1", None),
-        (True, {"NEEDRESTART_SUSPEND": "original"}, None, "original"),
-        (False, {"NEEDRESTART_SUSPEND": "original"}, "1", "original"),
+        (True, {containerd.NEEDRESTART_SUSPEND: "original"}, None, "original"),
+        (False, {containerd.NEEDRESTART_SUSPEND: "original"}, "1", "original"),
     ],
     ids=[
         "restart=True, no initial env",
@@ -430,19 +430,20 @@ def test_apt_restart_services(mock_log, restart, initial_env, expected_during, e
     """Verify _apt_restart_services behavior with various configurations."""
     # Setup initial environment from parameters
     os.environ.update(initial_env)
+    env_var = containerd.NEEDRESTART_SUSPEND
 
     with containerd._apt_restart_services(restart=restart):
         # Check value during context
         if expected_during is None:
-            assert "NEEDRESTART_SUSPEND" not in os.environ
+            assert env_var not in os.environ
         else:
-            assert os.environ["NEEDRESTART_SUSPEND"] == expected_during
+            assert os.environ[env_var] == expected_during
 
     # Check value after context
     if expected_after is None:
-        assert "NEEDRESTART_SUSPEND" not in os.environ
+        assert env_var not in os.environ
     else:
-        assert os.environ["NEEDRESTART_SUSPEND"] == expected_after
+        assert os.environ[env_var] == expected_after
 
     # Check log call
     log_fmt = "" if restart else "not "