diff --git a/src/local_ops.py b/src/local_ops.py index 1278689..db74224 100644 --- a/src/local_ops.py +++ b/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 diff --git a/src/os_ops.py b/src/os_ops.py index 6c9c453..35b29df 100644 --- a/src/os_ops.py +++ b/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 diff --git a/src/remote_ops.py b/src/remote_ops.py index 38b30cf..8d2424f 100644 --- a/src/remote_ops.py +++ b/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" diff --git a/tests/test_os_ops_common.py b/tests/test_os_ops_common.py index 97aa46c..bcfb6a1 100644 --- a/tests/test_os_ops_common.py +++ b/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: