new: OsOperations::reset_env is added

src/local_ops.py

@@ -414,10 +414,19 @@ def set_env(
         assert var_val is None or type(var_val) is str
         assert var_name != ""
 
-        if var_val is None:
-            os.environ.pop(var_name, None)
-        else:
-            os.environ[var_name] = var_val
+        self._set_env(var_name, var_val)
+        return
+
+    def reset_env(
+        self,
+        var_name: str,
+        default_val: typing.Optional[str],
+    ) -> None:
+        assert type(var_name) is str
+        assert default_val is None or type(default_val) is str
+        assert var_name != ""
+
+        self._set_env(var_name, default_val)
         return
 
     def get_name(self):
@@ -872,3 +881,18 @@ def _join_command_arguments(cmd: typing.Iterable[str]) -> str:
             assert type(item) is str
 
         return " ".join(__class__._quote_path(arg) for arg in cmd)
+
+    def _set_env(
+        self,
+        var_name: str,
+        var_val: typing.Optional[str],
+    ) -> None:
+        assert type(var_name) is str
+        assert var_val is None or type(var_val) is str
+        assert var_name != ""
+
+        if var_val is None:
+            os.environ.pop(var_name, None)
+        else:
+            os.environ[var_name] = var_val
+        return

src/os_ops.py

@@ -155,6 +155,16 @@ def set_env(
         assert var_name != ""
         raise NotImplementedError()
 
+    def reset_env(
+        self,
+        var_name: str,
+        default_val: typing.Optional[str],
+    ) -> None:
+        assert type(var_name) is str
+        assert default_val is None or type(default_val) is str
+        assert var_name != ""
+        raise NotImplementedError()
+
     def get_user(self) -> typing.Optional[str]:
         return self.username
 

src/remote_ops.py

@@ -430,6 +430,23 @@ def set_env(
             self._remote_env[var_name] = var_val
         return
 
+    def reset_env(
+        self,
+        var_name: str,
+        default_val: typing.Optional[str],
+    ) -> None:
+        assert type(var_name) is str
+        assert default_val is None or type(default_val) is str
+        assert var_name != ""
+
+        assert self._remote_env_guard is not None
+        assert type(self._remote_env) is dict
+
+        with self._remote_env_guard:
+            self._remote_env.pop(var_name, None)
+
+        return
+
     def get_name(self):
         return "posix"
 

tests/test_os_ops_common.py

@@ -220,22 +220,37 @@ def test_create_clone(
         os_ops = __class__.helper__get_os_ops(use_clone, os_ops_descr)
         assert isinstance(os_ops, OsOperations)
 
-        clone = os_ops.create_clone()
-        assert clone is not None
-        assert clone is not os_ops
-        assert type(clone) is type(os_ops)
+        env1_name = "env1_" + uuid.uuid4().bytes.hex()
+        env1_val = "abc"
+        env2_name = "env1_" + uuid.uuid4().bytes.hex()
 
-        assert clone.remote == os_ops.remote
-        assert clone.username == os_ops.username
-        assert clone.ssh_key == os_ops.ssh_key
-        assert clone.host == os_ops.host
-        assert clone.port == os_ops.port
+        os_ops.set_env(env1_name, env1_val)
 
-        assert clone.get_name() == os_ops.get_name()
-        assert clone.get_platform() == os_ops.get_platform()
+        try:
+            clone = os_ops.create_clone()
+            assert clone is not None
+            assert clone is not os_ops
+            assert type(clone) is type(os_ops)
 
-        assert clone.environ("PATH") == os_ops.environ("PATH")
-        assert clone.environ("DUMMY-4231") == os_ops.environ("DUMMY-4231")
+            assert clone.remote == os_ops.remote
+            assert clone.username == os_ops.username
+            assert clone.ssh_key == os_ops.ssh_key
+            assert clone.host == os_ops.host
+            assert clone.port == os_ops.port
+
+            assert clone.get_name() == os_ops.get_name()
+            assert clone.get_platform() == os_ops.get_platform()
+
+            assert clone.environ("PATH") == os_ops.environ("PATH")
+
+            v1_orig = os_ops.environ(env1_name)
+            v1_clone = clone.environ(env1_name)
+            assert v1_orig == env1_val
+            assert v1_orig == v1_clone
+
+            assert clone.environ(env2_name) == os_ops.environ(env2_name)
+        finally:
+            os_ops.reset_env(env1_name, None)
 
         return
 
@@ -3547,61 +3562,70 @@ def test_environ(
     def test_set_env__persistence(
         self,
         os_ops_descr: OsOpsDescr,
+        use_clone: bool,
     ):
         assert type(os_ops_descr) is OsOpsDescr
-        os_ops = os_ops_descr.os_ops
+        assert type(use_clone) is bool
+
+        os_ops = __class__.helper__get_os_ops(use_clone, os_ops_descr)
         assert isinstance(os_ops, OsOperations)
 
         var_name = "TEST_SET_ENV_MAGIC_VAR"
-        var_value = "Cokanum_Detected_123"
 
-        # Step 1: Set the variable
-        os_ops.set_env(var_name, var_value)
+        assert os_ops.environ(var_name) is None
 
-        # Step 2: In a SEPARATE command, we try to read it using our new environ()
-        # The old remote_ops is guaranteed to return None and crash!
-        fetched_value = os_ops.environ(var_name)
+        try:
+            var_value = "Cokanum_Detected_123"
 
-        assert fetched_value == var_value, "Env variable persistence failed! Got: {}".format(fetched_value)
-        logging.info("SUCCESS. Environment variable persistence verified across commands.")
+            # Step 1: Set the variable
+            os_ops.set_env(var_name, var_value)
 
-        printenv = os_ops.find_executable("printenv")
-        assert type(printenv) is str
-        assert printenv != ""
+            # Step 2: In a SEPARATE command, we try to read it using our new environ()
+            # The old remote_ops is guaranteed to return None and crash!
+            fetched_value = os_ops.environ(var_name)
 
-        cmd1 = [printenv, var_name]
+            assert fetched_value == var_value, "Env variable persistence failed! Got: {}".format(fetched_value)
+            logging.info("SUCCESS. Environment variable persistence verified across commands.")
 
-        exec_r = os_ops.exec_command(
-            cmd1,
-            encoding="utf-8",
-        )
-        assert type(exec_r) is str
-        exec_r = exec_r.rstrip()
-        assert fetched_value == var_value
+            printenv = os_ops.find_executable("printenv")
+            assert type(printenv) is str
+            assert printenv != ""
 
-        exec_r = os_ops.exec_command(
-            cmd1,
-            encoding="utf-8",
-            exec_env={
-                var_name: "ABC",
-            }
-        )
-        assert type(exec_r) is str
-        exec_r = exec_r.rstrip()
-        assert exec_r == "ABC"
+            cmd1 = [printenv, var_name]
 
-        exec_r = os_ops.exec_command(
-            cmd1,
-            encoding="utf-8",
-            exec_env={
-                var_name: None,
-            },
-            ignore_errors=True,
-            verbose=True,
-        )
-        assert type(exec_r) is tuple
-        assert len(exec_r) == 3
-        assert exec_r[0] == 1
+            exec_r = os_ops.exec_command(
+                cmd1,
+                encoding="utf-8",
+            )
+            assert type(exec_r) is str
+            exec_r = exec_r.rstrip()
+            assert fetched_value == var_value
+
+            exec_r = os_ops.exec_command(
+                cmd1,
+                encoding="utf-8",
+                exec_env={
+                    var_name: "ABC",
+                }
+            )
+            assert type(exec_r) is str
+            exec_r = exec_r.rstrip()
+            assert exec_r == "ABC"
+
+            exec_r = os_ops.exec_command(
+                cmd1,
+                encoding="utf-8",
+                exec_env={
+                    var_name: None,
+                },
+                ignore_errors=True,
+                verbose=True,
+            )
+            assert type(exec_r) is tuple
+            assert len(exec_r) == 3
+            assert exec_r[0] == 1
+        finally:
+            os_ops.reset_env(var_name, None)
 
         # ----------------
         x = os_ops.environ("PATH")
@@ -3621,11 +3645,41 @@ def test_set_env__persistence(
             assert len(exec_r) == 3
             assert exec_r[0] == 1
         finally:
-            os_ops.set_env("PATH", x)
+            os_ops.reset_env("PATH", x)
             assert os_ops.environ("PATH") == x
 
         return
 
+    def test_reset_env(
+        self,
+        os_ops_descr: OsOpsDescr,
+        use_clone: bool,
+    ):
+        assert type(os_ops_descr) is OsOpsDescr
+        assert type(use_clone) is bool
+
+        os_ops = __class__.helper__get_os_ops(use_clone, os_ops_descr)
+        assert isinstance(os_ops, OsOperations)
+
+        C_VAR_NAME = "PATH"
+
+        clone = __class__.helper__create_clone_and_formal_check_it(os_ops)
+
+        origin = os_ops.environ(C_VAR_NAME)
+        assert origin is not None
+        assert type(origin) is str
+        newvalue = origin + os_ops.pathsep + "aaaa"
+
+        os_ops.set_env(C_VAR_NAME, newvalue)
+
+        assert os_ops.environ(C_VAR_NAME) == newvalue
+
+        os_ops.reset_env(C_VAR_NAME, origin)
+
+        assert os_ops.environ(C_VAR_NAME) == origin
+        assert clone.environ(C_VAR_NAME) == origin
+        return
+
     def test_set_env__evil(
         self,
         os_ops_descr: OsOpsDescr,
@@ -3699,7 +3753,7 @@ def thread_worker(var_name: str, var_value: str, iterations: int):
                         )
 
                     # 3. Clean up after yourself
-                    os_ops.set_env(var_name, None)
+                    os_ops.reset_env(var_name, None)
                     assert os_ops.environ(var_name) is None
 
             except Exception as e: