Skip to content

Commit a86d1c5

Browse files
new: OsOperations::reset_env is added (#108)
Description: def reset_env( self, var_name: str, default_val: typing.Optional[str], ) -> None: LocalOperations::reset_env is equal to set_env method: - If default_val is not None, it is assigned into a variable. - If default_val is None, a variable is deleted. RemoteOperations::reset_env removes a variable from own local cache. Parameter default_val is ignored. --- Updated tests: - test_create_clone - test_set_env__persistence - test_set_env__thread_safety
1 parent 1855033 commit a86d1c5

4 files changed

Lines changed: 166 additions & 61 deletions

File tree

src/local_ops.py

Lines changed: 28 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -414,10 +414,19 @@ def set_env(
414414
assert var_val is None or type(var_val) is str
415415
assert var_name != ""
416416

417-
if var_val is None:
418-
os.environ.pop(var_name, None)
419-
else:
420-
os.environ[var_name] = var_val
417+
self._set_env(var_name, var_val)
418+
return
419+
420+
def reset_env(
421+
self,
422+
var_name: str,
423+
default_val: typing.Optional[str],
424+
) -> None:
425+
assert type(var_name) is str
426+
assert default_val is None or type(default_val) is str
427+
assert var_name != ""
428+
429+
self._set_env(var_name, default_val)
421430
return
422431

423432
def get_name(self):
@@ -872,3 +881,18 @@ def _join_command_arguments(cmd: typing.Iterable[str]) -> str:
872881
assert type(item) is str
873882

874883
return " ".join(__class__._quote_path(arg) for arg in cmd)
884+
885+
def _set_env(
886+
self,
887+
var_name: str,
888+
var_val: typing.Optional[str],
889+
) -> None:
890+
assert type(var_name) is str
891+
assert var_val is None or type(var_val) is str
892+
assert var_name != ""
893+
894+
if var_val is None:
895+
os.environ.pop(var_name, None)
896+
else:
897+
os.environ[var_name] = var_val
898+
return

src/os_ops.py

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -155,6 +155,16 @@ def set_env(
155155
assert var_name != ""
156156
raise NotImplementedError()
157157

158+
def reset_env(
159+
self,
160+
var_name: str,
161+
default_val: typing.Optional[str],
162+
) -> None:
163+
assert type(var_name) is str
164+
assert default_val is None or type(default_val) is str
165+
assert var_name != ""
166+
raise NotImplementedError()
167+
158168
def get_user(self) -> typing.Optional[str]:
159169
return self.username
160170

src/remote_ops.py

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -430,6 +430,23 @@ def set_env(
430430
self._remote_env[var_name] = var_val
431431
return
432432

433+
def reset_env(
434+
self,
435+
var_name: str,
436+
default_val: typing.Optional[str],
437+
) -> None:
438+
assert type(var_name) is str
439+
assert default_val is None or type(default_val) is str
440+
assert var_name != ""
441+
442+
assert self._remote_env_guard is not None
443+
assert type(self._remote_env) is dict
444+
445+
with self._remote_env_guard:
446+
self._remote_env.pop(var_name, None)
447+
448+
return
449+
433450
def get_name(self):
434451
return "posix"
435452

tests/test_os_ops_common.py

Lines changed: 111 additions & 57 deletions
Original file line numberDiff line numberDiff line change
@@ -220,22 +220,37 @@ def test_create_clone(
220220
os_ops = __class__.helper__get_os_ops(use_clone, os_ops_descr)
221221
assert isinstance(os_ops, OsOperations)
222222

223-
clone = os_ops.create_clone()
224-
assert clone is not None
225-
assert clone is not os_ops
226-
assert type(clone) is type(os_ops)
223+
env1_name = "env1_" + uuid.uuid4().bytes.hex()
224+
env1_val = "abc"
225+
env2_name = "env1_" + uuid.uuid4().bytes.hex()
227226

228-
assert clone.remote == os_ops.remote
229-
assert clone.username == os_ops.username
230-
assert clone.ssh_key == os_ops.ssh_key
231-
assert clone.host == os_ops.host
232-
assert clone.port == os_ops.port
227+
os_ops.set_env(env1_name, env1_val)
233228

234-
assert clone.get_name() == os_ops.get_name()
235-
assert clone.get_platform() == os_ops.get_platform()
229+
try:
230+
clone = os_ops.create_clone()
231+
assert clone is not None
232+
assert clone is not os_ops
233+
assert type(clone) is type(os_ops)
236234

237-
assert clone.environ("PATH") == os_ops.environ("PATH")
238-
assert clone.environ("DUMMY-4231") == os_ops.environ("DUMMY-4231")
235+
assert clone.remote == os_ops.remote
236+
assert clone.username == os_ops.username
237+
assert clone.ssh_key == os_ops.ssh_key
238+
assert clone.host == os_ops.host
239+
assert clone.port == os_ops.port
240+
241+
assert clone.get_name() == os_ops.get_name()
242+
assert clone.get_platform() == os_ops.get_platform()
243+
244+
assert clone.environ("PATH") == os_ops.environ("PATH")
245+
246+
v1_orig = os_ops.environ(env1_name)
247+
v1_clone = clone.environ(env1_name)
248+
assert v1_orig == env1_val
249+
assert v1_orig == v1_clone
250+
251+
assert clone.environ(env2_name) == os_ops.environ(env2_name)
252+
finally:
253+
os_ops.reset_env(env1_name, None)
239254

240255
return
241256

@@ -3547,61 +3562,70 @@ def test_environ(
35473562
def test_set_env__persistence(
35483563
self,
35493564
os_ops_descr: OsOpsDescr,
3565+
use_clone: bool,
35503566
):
35513567
assert type(os_ops_descr) is OsOpsDescr
3552-
os_ops = os_ops_descr.os_ops
3568+
assert type(use_clone) is bool
3569+
3570+
os_ops = __class__.helper__get_os_ops(use_clone, os_ops_descr)
35533571
assert isinstance(os_ops, OsOperations)
35543572

35553573
var_name = "TEST_SET_ENV_MAGIC_VAR"
3556-
var_value = "Cokanum_Detected_123"
35573574

3558-
# Step 1: Set the variable
3559-
os_ops.set_env(var_name, var_value)
3575+
assert os_ops.environ(var_name) is None
35603576

3561-
# Step 2: In a SEPARATE command, we try to read it using our new environ()
3562-
# The old remote_ops is guaranteed to return None and crash!
3563-
fetched_value = os_ops.environ(var_name)
3577+
try:
3578+
var_value = "Cokanum_Detected_123"
35643579

3565-
assert fetched_value == var_value, "Env variable persistence failed! Got: {}".format(fetched_value)
3566-
logging.info("SUCCESS. Environment variable persistence verified across commands.")
3580+
# Step 1: Set the variable
3581+
os_ops.set_env(var_name, var_value)
35673582

3568-
printenv = os_ops.find_executable("printenv")
3569-
assert type(printenv) is str
3570-
assert printenv != ""
3583+
# Step 2: In a SEPARATE command, we try to read it using our new environ()
3584+
# The old remote_ops is guaranteed to return None and crash!
3585+
fetched_value = os_ops.environ(var_name)
35713586

3572-
cmd1 = [printenv, var_name]
3587+
assert fetched_value == var_value, "Env variable persistence failed! Got: {}".format(fetched_value)
3588+
logging.info("SUCCESS. Environment variable persistence verified across commands.")
35733589

3574-
exec_r = os_ops.exec_command(
3575-
cmd1,
3576-
encoding="utf-8",
3577-
)
3578-
assert type(exec_r) is str
3579-
exec_r = exec_r.rstrip()
3580-
assert fetched_value == var_value
3590+
printenv = os_ops.find_executable("printenv")
3591+
assert type(printenv) is str
3592+
assert printenv != ""
35813593

3582-
exec_r = os_ops.exec_command(
3583-
cmd1,
3584-
encoding="utf-8",
3585-
exec_env={
3586-
var_name: "ABC",
3587-
}
3588-
)
3589-
assert type(exec_r) is str
3590-
exec_r = exec_r.rstrip()
3591-
assert exec_r == "ABC"
3594+
cmd1 = [printenv, var_name]
35923595

3593-
exec_r = os_ops.exec_command(
3594-
cmd1,
3595-
encoding="utf-8",
3596-
exec_env={
3597-
var_name: None,
3598-
},
3599-
ignore_errors=True,
3600-
verbose=True,
3601-
)
3602-
assert type(exec_r) is tuple
3603-
assert len(exec_r) == 3
3604-
assert exec_r[0] == 1
3596+
exec_r = os_ops.exec_command(
3597+
cmd1,
3598+
encoding="utf-8",
3599+
)
3600+
assert type(exec_r) is str
3601+
exec_r = exec_r.rstrip()
3602+
assert fetched_value == var_value
3603+
3604+
exec_r = os_ops.exec_command(
3605+
cmd1,
3606+
encoding="utf-8",
3607+
exec_env={
3608+
var_name: "ABC",
3609+
}
3610+
)
3611+
assert type(exec_r) is str
3612+
exec_r = exec_r.rstrip()
3613+
assert exec_r == "ABC"
3614+
3615+
exec_r = os_ops.exec_command(
3616+
cmd1,
3617+
encoding="utf-8",
3618+
exec_env={
3619+
var_name: None,
3620+
},
3621+
ignore_errors=True,
3622+
verbose=True,
3623+
)
3624+
assert type(exec_r) is tuple
3625+
assert len(exec_r) == 3
3626+
assert exec_r[0] == 1
3627+
finally:
3628+
os_ops.reset_env(var_name, None)
36053629

36063630
# ----------------
36073631
x = os_ops.environ("PATH")
@@ -3621,11 +3645,41 @@ def test_set_env__persistence(
36213645
assert len(exec_r) == 3
36223646
assert exec_r[0] == 1
36233647
finally:
3624-
os_ops.set_env("PATH", x)
3648+
os_ops.reset_env("PATH", x)
36253649
assert os_ops.environ("PATH") == x
36263650

36273651
return
36283652

3653+
def test_reset_env(
3654+
self,
3655+
os_ops_descr: OsOpsDescr,
3656+
use_clone: bool,
3657+
):
3658+
assert type(os_ops_descr) is OsOpsDescr
3659+
assert type(use_clone) is bool
3660+
3661+
os_ops = __class__.helper__get_os_ops(use_clone, os_ops_descr)
3662+
assert isinstance(os_ops, OsOperations)
3663+
3664+
C_VAR_NAME = "PATH"
3665+
3666+
clone = __class__.helper__create_clone_and_formal_check_it(os_ops)
3667+
3668+
origin = os_ops.environ(C_VAR_NAME)
3669+
assert origin is not None
3670+
assert type(origin) is str
3671+
newvalue = origin + os_ops.pathsep + "aaaa"
3672+
3673+
os_ops.set_env(C_VAR_NAME, newvalue)
3674+
3675+
assert os_ops.environ(C_VAR_NAME) == newvalue
3676+
3677+
os_ops.reset_env(C_VAR_NAME, origin)
3678+
3679+
assert os_ops.environ(C_VAR_NAME) == origin
3680+
assert clone.environ(C_VAR_NAME) == origin
3681+
return
3682+
36293683
def test_set_env__evil(
36303684
self,
36313685
os_ops_descr: OsOpsDescr,
@@ -3699,7 +3753,7 @@ def thread_worker(var_name: str, var_value: str, iterations: int):
36993753
)
37003754

37013755
# 3. Clean up after yourself
3702-
os_ops.set_env(var_name, None)
3756+
os_ops.reset_env(var_name, None)
37033757
assert os_ops.environ(var_name) is None
37043758

37053759
except Exception as e:

0 commit comments

Comments
 (0)