Skip to content

Commit 3575e2e

Browse files
godobytegodobyte
committed
feat: optionally apply existing env from session context for run_task
Signed-off-by: Godot Bian <13778003+godobyte@users.noreply.github.com>
1 parent 5fc31f0 commit 3575e2e

File tree

2 files changed

+295
-7
lines changed

2 files changed

+295
-7
lines changed

‎src/openjd/sessions/_session.py‎

Lines changed: 20 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -782,6 +782,7 @@ def run_task(
782782
step_script: StepScriptModel,
783783
task_parameter_values: TaskParameterSet,
784784
os_env_vars: Optional[dict[str, str]] = None,
785+
apply_existing_env: Optional[bool] = True,
785786
log_task_banner: bool = True,
786787
) -> None:
787788
"""Run a Task within the Session.
@@ -799,6 +800,9 @@ def run_task(
799800
by values defined in Environments.
800801
Key: Environment variable name
801802
Value: Value for the environment variable.
803+
apply_existing_env (Optional[bool]): Whether to apply environment variable changes from
804+
entered environments. When True (default), applies changes from all entered environments
805+
in order. When False, only uses base session and os_env_vars.
802806
log_task_banner (bool): Whether to log a banner before running the Task.
803807
Default: True
804808
"""
@@ -821,7 +825,7 @@ def run_task(
821825

822826
self._reset_action_state()
823827
symtab = self._symbol_table(step_script.revision, task_parameter_values)
824-
action_env_vars = self._evaluate_current_session_env_vars(os_env_vars)
828+
action_env_vars = self._evaluate_current_session_env_vars(os_env_vars, apply_existing_env)
825829
self._materialize_path_mapping(step_script.revision, action_env_vars, symtab)
826830
self._runner = StepScriptRunner(
827831
logger=self._logger,
@@ -1120,15 +1124,24 @@ def _action_callback(self, state: ActionState) -> None:
11201124
self._callback(self._session_id, action_status)
11211125

11221126
def _evaluate_current_session_env_vars(
1123-
self, extra_env_vars: Optional[dict[str, str]] = None
1127+
self,
1128+
extra_env_vars: Optional[dict[str, str]] = None,
1129+
apply_existing_env: Optional[bool] = True,
11241130
) -> dict[str, Optional[str]]:
1125-
"""Get a dictionary representing the cummulative state of env vars set
1126-
and unset from the currently applied environments.
1131+
"""Get a dictionary
1132+
1. representing the cummulative state of env vars set and unset
1133+
from the currently applied environments.
1134+
2. when apply_existing_env is provided and set to False,
1135+
skip applying the current environments.
11271136
"""
11281137
result = dict[str, Optional[str]](self._process_env) # Make a copy
1138+
11291139
if extra_env_vars:
11301140
result.update(**extra_env_vars)
1131-
for identifier in self._environments_entered:
1132-
if identifier in self._created_env_vars:
1133-
self._created_env_vars[identifier].apply_to_environment(result)
1141+
1142+
if apply_existing_env:
1143+
for identifier in self._environments_entered:
1144+
if identifier in self._created_env_vars:
1145+
self._created_env_vars[identifier].apply_to_environment(result)
1146+
11341147
return result

‎test/openjd/sessions/test_session.py‎

Lines changed: 275 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -929,6 +929,281 @@ def test_run_task_with_variables(
929929
assert fix_foo_baz_environment.variables is not None
930930
assert session._runner._os_env_vars == dict(fix_foo_baz_environment.variables)
931931

932+
def test_run_task_with_apply_existing_env_true(
933+
self,
934+
caplog: pytest.LogCaptureFixture,
935+
python_exe: str,
936+
fix_foo_baz_environment: Environment_2023_09,
937+
) -> None:
938+
# Test that when apply_existing_env=True (default), environment variables from entered environments are applied
939+
# GIVEN
940+
step_script = StepScript_2023_09(
941+
actions=StepActions_2023_09(
942+
onRun=Action_2023_09(
943+
command=CommandString_2023_09(python_exe),
944+
args=[ArgString_2023_09("{{ Task.File.Foo }}")],
945+
)
946+
),
947+
embeddedFiles=[
948+
EmbeddedFileText_2023_09(
949+
name="Foo",
950+
type=EmbeddedFileTypes_2023_09.TEXT,
951+
data=DataString_2023_09(
952+
'import time; import os; time.sleep(0.5); print(f\'{os.environ.get("FOO", "NOT_SET")} {os.environ.get("BAZ", "NOT_SET")}\')'
953+
),
954+
)
955+
],
956+
)
957+
958+
session_id = uuid.uuid4().hex
959+
job_params = dict[str, ParameterValue]()
960+
task_params = dict[str, ParameterValue]()
961+
962+
with Session(session_id=session_id, job_parameter_values=job_params) as session:
963+
# Enter an environment that sets FOO and BAZ
964+
session.enter_environment(environment=fix_foo_baz_environment)
965+
966+
# WHEN - apply_existing_env=True (default)
967+
session.run_task(
968+
step_script=step_script,
969+
task_parameter_values=task_params,
970+
apply_existing_env=True,
971+
)
972+
973+
# THEN
974+
assert session.state == SessionState.RUNNING
975+
# Wait for the process to exit
976+
while session.state == SessionState.RUNNING:
977+
time.sleep(0.1)
978+
assert session.state == SessionState.READY
979+
assert session.action_status == ActionStatus(state=ActionState.SUCCESS, exit_code=0)
980+
# Environment variables from entered environment should be available
981+
assert "bar qux" in caplog.messages
982+
983+
def test_run_task_with_apply_existing_env_false(
984+
self,
985+
caplog: pytest.LogCaptureFixture,
986+
python_exe: str,
987+
fix_foo_baz_environment: Environment_2023_09,
988+
) -> None:
989+
# Test that when apply_existing_env=False, environment variables from entered environments are NOT applied
990+
# GIVEN
991+
step_script = StepScript_2023_09(
992+
actions=StepActions_2023_09(
993+
onRun=Action_2023_09(
994+
command=CommandString_2023_09(python_exe),
995+
args=[ArgString_2023_09("{{ Task.File.Foo }}")],
996+
)
997+
),
998+
embeddedFiles=[
999+
EmbeddedFileText_2023_09(
1000+
name="Foo",
1001+
type=EmbeddedFileTypes_2023_09.TEXT,
1002+
data=DataString_2023_09(
1003+
'import time; import os; time.sleep(0.5); print(f\'{os.environ.get("FOO", "NOT_SET")} {os.environ.get("BAZ", "NOT_SET")}\')'
1004+
),
1005+
)
1006+
],
1007+
)
1008+
1009+
session_id = uuid.uuid4().hex
1010+
job_params = dict[str, ParameterValue]()
1011+
task_params = dict[str, ParameterValue]()
1012+
1013+
with Session(session_id=session_id, job_parameter_values=job_params) as session:
1014+
# Enter an environment that sets FOO and BAZ
1015+
session.enter_environment(environment=fix_foo_baz_environment)
1016+
1017+
# WHEN - apply_existing_env=False
1018+
session.run_task(
1019+
step_script=step_script,
1020+
task_parameter_values=task_params,
1021+
apply_existing_env=False,
1022+
)
1023+
1024+
# THEN
1025+
assert session.state == SessionState.RUNNING
1026+
# Wait for the process to exit
1027+
while session.state == SessionState.RUNNING:
1028+
time.sleep(0.1)
1029+
assert session.state == SessionState.READY
1030+
assert session.action_status == ActionStatus(state=ActionState.SUCCESS, exit_code=0)
1031+
# Environment variables from entered environment should NOT be available
1032+
assert "NOT_SET NOT_SET" in caplog.messages
1033+
1034+
def test_run_task_apply_existing_env_with_extra_env_vars(
1035+
self,
1036+
caplog: pytest.LogCaptureFixture,
1037+
python_exe: str,
1038+
fix_foo_baz_environment: Environment_2023_09,
1039+
) -> None:
1040+
# Test that extra_env_vars are always applied regardless of apply_existing_env setting
1041+
# GIVEN
1042+
step_script = StepScript_2023_09(
1043+
actions=StepActions_2023_09(
1044+
onRun=Action_2023_09(
1045+
command=CommandString_2023_09(python_exe),
1046+
args=[ArgString_2023_09("{{ Task.File.Foo }}")],
1047+
)
1048+
),
1049+
embeddedFiles=[
1050+
EmbeddedFileText_2023_09(
1051+
name="Foo",
1052+
type=EmbeddedFileTypes_2023_09.TEXT,
1053+
data=DataString_2023_09(
1054+
'import time; import os; time.sleep(0.5); print(f\'{os.environ.get("FOO", "NOT_SET")} {os.environ.get("EXTRA_VAR", "NOT_SET")}\')'
1055+
),
1056+
)
1057+
],
1058+
)
1059+
1060+
session_id = uuid.uuid4().hex
1061+
job_params = dict[str, ParameterValue]()
1062+
task_params = dict[str, ParameterValue]()
1063+
extra_env_vars = {"EXTRA_VAR": "extra_value"}
1064+
1065+
with Session(session_id=session_id, job_parameter_values=job_params) as session:
1066+
# Enter an environment that sets FOO and BAZ
1067+
session.enter_environment(environment=fix_foo_baz_environment)
1068+
1069+
# WHEN - apply_existing_env=False but with extra_env_vars
1070+
session.run_task(
1071+
step_script=step_script,
1072+
task_parameter_values=task_params,
1073+
os_env_vars=extra_env_vars,
1074+
apply_existing_env=False,
1075+
)
1076+
1077+
# THEN
1078+
assert session.state == SessionState.RUNNING
1079+
# Wait for the process to exit
1080+
while session.state == SessionState.RUNNING:
1081+
time.sleep(0.1)
1082+
assert session.state == SessionState.READY
1083+
assert session.action_status == ActionStatus(state=ActionState.SUCCESS, exit_code=0)
1084+
# Environment variables from entered environment should NOT be available, but extra_env_vars should be
1085+
assert "NOT_SET extra_value" in caplog.messages
1086+
1087+
def test_run_task_apply_existing_env_default_behavior(
1088+
self,
1089+
caplog: pytest.LogCaptureFixture,
1090+
python_exe: str,
1091+
fix_foo_baz_environment: Environment_2023_09,
1092+
) -> None:
1093+
# Test that the default behavior (when apply_existing_env is not specified) is to apply existing env vars
1094+
# GIVEN
1095+
step_script = StepScript_2023_09(
1096+
actions=StepActions_2023_09(
1097+
onRun=Action_2023_09(
1098+
command=CommandString_2023_09(python_exe),
1099+
args=[ArgString_2023_09("{{ Task.File.Foo }}")],
1100+
)
1101+
),
1102+
embeddedFiles=[
1103+
EmbeddedFileText_2023_09(
1104+
name="Foo",
1105+
type=EmbeddedFileTypes_2023_09.TEXT,
1106+
data=DataString_2023_09(
1107+
'import time; import os; time.sleep(0.5); print(f\'{os.environ.get("FOO", "NOT_SET")} {os.environ.get("BAZ", "NOT_SET")}\')'
1108+
),
1109+
)
1110+
],
1111+
)
1112+
1113+
session_id = uuid.uuid4().hex
1114+
job_params = dict[str, ParameterValue]()
1115+
task_params = dict[str, ParameterValue]()
1116+
1117+
with Session(session_id=session_id, job_parameter_values=job_params) as session:
1118+
# Enter an environment that sets FOO and BAZ
1119+
session.enter_environment(environment=fix_foo_baz_environment)
1120+
1121+
# WHEN - apply_existing_env not specified (should default to True)
1122+
session.run_task(
1123+
step_script=step_script,
1124+
task_parameter_values=task_params,
1125+
)
1126+
1127+
# THEN
1128+
assert session.state == SessionState.RUNNING
1129+
# Wait for the process to exit
1130+
while session.state == SessionState.RUNNING:
1131+
time.sleep(0.1)
1132+
assert session.state == SessionState.READY
1133+
assert session.action_status == ActionStatus(state=ActionState.SUCCESS, exit_code=0)
1134+
# Environment variables from entered environment should be available (default behavior)
1135+
assert "bar qux" in caplog.messages
1136+
1137+
def test_evaluate_current_session_env_vars_with_apply_existing_env_true(self) -> None:
1138+
# Test that when apply_existing_env=True, environment variables from entered environments are included
1139+
# GIVEN
1140+
session_id = uuid.uuid4().hex
1141+
job_params = dict[str, ParameterValue]()
1142+
1143+
with Session(session_id=session_id, job_parameter_values=job_params) as session:
1144+
# Set up some environment variables in the session
1145+
session._created_env_vars["test_env"] = SimplifiedEnvironmentVariableChanges(
1146+
{"FOO": "bar", "BAZ": "qux"}
1147+
)
1148+
session._environments_entered.append("test_env")
1149+
1150+
extra_env_vars = {"EXTRA": "value"}
1151+
1152+
# WHEN
1153+
result = session._evaluate_current_session_env_vars(
1154+
extra_env_vars=extra_env_vars, apply_existing_env=True
1155+
)
1156+
1157+
# THEN
1158+
assert result["EXTRA"] == "value" # Extra env vars should be included
1159+
assert result["FOO"] == "bar" # Environment vars should be included
1160+
assert result["BAZ"] == "qux" # Environment vars should be included
1161+
1162+
def test_evaluate_current_session_env_vars_with_apply_existing_env_false(self) -> None:
1163+
# Test that when apply_existing_env=False, environment variables from entered environments are NOT included
1164+
# GIVEN
1165+
session_id = uuid.uuid4().hex
1166+
job_params = dict[str, ParameterValue]()
1167+
1168+
with Session(session_id=session_id, job_parameter_values=job_params) as session:
1169+
# Set up some environment variables in the session
1170+
session._created_env_vars["test_env"] = SimplifiedEnvironmentVariableChanges(
1171+
{"FOO": "bar", "BAZ": "qux"}
1172+
)
1173+
session._environments_entered.append("test_env")
1174+
1175+
extra_env_vars = {"EXTRA": "value"}
1176+
1177+
# WHEN
1178+
result = session._evaluate_current_session_env_vars(
1179+
extra_env_vars=extra_env_vars, apply_existing_env=False
1180+
)
1181+
1182+
# THEN
1183+
assert result["EXTRA"] == "value" # Extra env vars should be included
1184+
assert "FOO" not in result # Environment vars should NOT be included
1185+
assert "BAZ" not in result # Environment vars should NOT be included
1186+
1187+
def test_evaluate_current_session_env_vars_default_apply_existing_env(self) -> None:
1188+
# Test that the default behavior (apply_existing_env=True) includes environment variables
1189+
# GIVEN
1190+
session_id = uuid.uuid4().hex
1191+
job_params = dict[str, ParameterValue]()
1192+
1193+
with Session(session_id=session_id, job_parameter_values=job_params) as session:
1194+
# Set up some environment variables in the session
1195+
session._created_env_vars["test_env"] = SimplifiedEnvironmentVariableChanges(
1196+
{"FOO": "bar", "BAZ": "qux"}
1197+
)
1198+
session._environments_entered.append("test_env")
1199+
1200+
# WHEN - not specifying apply_existing_env (should default to True)
1201+
result = session._evaluate_current_session_env_vars()
1202+
1203+
# THEN
1204+
assert result["FOO"] == "bar" # Environment vars should be included by default
1205+
assert result["BAZ"] == "qux" # Environment vars should be included by default
1206+
9321207

9331208
class TestSessionCancel:
9341209
"""Test that cancelation will cancel the currently running Script."""

0 commit comments

Comments
 (0)