feat: Support for multi line env variables in enter env (#115)

Signed-off-by: Sakshi Sakshi <[email protected]>

src/openjd/sessions/_action_filter.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import json
 import logging
 import re
 from enum import Enum
@@ -36,8 +37,10 @@ class ActionMessageKind(Enum):
 openjd_env_actions_filter_matcher = re.compile(openjd_env_actions_filter_regex)
 
 # A regex for matching the assignment of a value to an environment variable
-envvar_set_regex = "^[A-Za-z_][A-Za-z0-9_]*" "=" ".*$"  # Variable name
-envvar_set_matcher = re.compile(envvar_set_regex)
+envvar_set_regex_str = "^[A-Za-z_][A-Za-z0-9_]*" "=" ".*$"  # Variable name
+envvar_set_regex_json = '^(")?[A-Za-z_][A-Za-z0-9_]*' "=" ".*$"  # Variable name
+envvar_set_matcher_str = re.compile(envvar_set_regex_str)
+envvar_set_matcher_json = re.compile(envvar_set_regex_json)
 envvar_unset_regex = "^[A-Za-z_][A-Za-z0-9_]*$"
 envvar_unset_matcher = re.compile(envvar_unset_regex)
 
@@ -233,12 +236,17 @@ def _handle_env(self, message: str) -> None:
         #   <varname> consists of latin alphanumeric characters and the underscore,
         #             and starts with a non-digit
         #   <value> can be any characters including empty.
-        if not envvar_set_matcher.match(message):
+        if not envvar_set_matcher_str.match(message) and not envvar_set_matcher_json.match(message):
             err_message = "Failed to parse environment variable assignment."
             # Callback to fail and cancel action on this error
             self._callback(ActionMessageKind.ENV, err_message, True)
             raise ValueError(err_message)
-        name, _, value = message.partition("=")
+        elif envvar_set_matcher_str.match(message):
+            name, _, value = message.partition("=")
+        else:
+            message_json_str = json.loads(message)
+            name, _, value = message_json_str.partition("=")
+
         self._callback(ActionMessageKind.ENV, {"name": name, "value": value}, False)
 
     def _handle_unset_env(self, message: str) -> None:

test/openjd/sessions/test_session.py

@@ -2358,6 +2358,89 @@ def test_def_via_stdout(
             assert "FOO=FOO-value" in caplog.messages
             assert "BAR=BAR-value" in caplog.messages
 
+    @pytest.mark.usefixtures("caplog")  # builtin fixture
+    def test_def_via_multi_line_nonvalid_json_stdout(
+        self, caplog: pytest.LogCaptureFixture, step_script_definition: StepScript_2023_09
+    ) -> None:
+        # Test that when an environment defines variables via a stdout handler
+        # as a multiline json it is processed properly
+
+        # GIVEN
+        environment = Environment_2023_09(
+            name="Env",
+            script=EnvironmentScript_2023_09(
+                actions=EnvironmentActions_2023_09(
+                    onEnter=Action_2023_09(
+                        command=sys.executable,
+                        args=["-c", "import json; print('openjd_env: \"FOO=12\\\\n34')"],
+                    )
+                )
+            ),
+        )
+        session_id = uuid.uuid4().hex
+        job_params = dict[str, ParameterValue]()
+        with Session(session_id=session_id, job_parameter_values=job_params) as session:
+            session.enter_environment(environment=environment)
+            while session.state == SessionState.RUNNING:
+                time.sleep(0.1)
+
+            assert (
+                'openjd_env: "FOO=12\\n34 -- ERROR: Unterminated string starting at: line 1 column 1 (char 0)'
+                in caplog.messages
+            )
+
+    @pytest.mark.usefixtures("caplog")  # builtin fixture
+    def test_def_via_multi_line_stdout(
+        self,
+        caplog: pytest.LogCaptureFixture,
+    ) -> None:
+        # Test that when an environment defines variables via a stdout handler
+        # as a multiline json it is processed properly
+
+        # GIVEN
+        environment = Environment_2023_09(
+            name="Env",
+            script=EnvironmentScript_2023_09(
+                actions=EnvironmentActions_2023_09(
+                    onEnter=Action_2023_09(
+                        command=sys.executable,
+                        args=["-c", "print('openjd_env: \"FOO=12\\\\n34\"')"],
+                    )
+                ),
+            ),
+        )
+
+        script = StepScript_2023_09(
+            actions=StepActions_2023_09(
+                onRun=Action_2023_09(
+                    command=sys.executable,
+                    args=[
+                        "-c",
+                        "import os; print('FOO:'); print(f'{os.environ[\"FOO\"]}'); print('---')",
+                    ],
+                )
+            ),
+        )
+        session_id = uuid.uuid4().hex
+        job_params = dict[str, ParameterValue]()
+        with Session(session_id=session_id, job_parameter_values=job_params) as session:
+            session.enter_environment(environment=environment)
+            while session.state == SessionState.RUNNING:
+                time.sleep(0.1)
+            # WHEN
+            session.run_task(
+                step_script=script,
+                task_parameter_values=dict[str, ParameterValue](),
+            )
+            while session.state == SessionState.RUNNING:
+                time.sleep(0.1)
+
+            # THEN
+            assert "FOO:" in caplog.messages
+            assert "12" in caplog.messages
+            assert "34" in caplog.messages
+            assert "---" in caplog.messages
+
     @pytest.mark.usefixtures("caplog")  # builtin fixture
     def test_def_via_stdout_overrides_direct(
         self, caplog: pytest.LogCaptureFixture, step_script_definition: StepScript_2023_09
@@ -2396,6 +2479,76 @@ def test_def_via_stdout_overrides_direct(
             assert "FOO=FOO-value" in caplog.messages
             assert "BAR=BAR-value" in caplog.messages
 
+    @pytest.mark.usefixtures("caplog")  # builtin fixture
+    def test_def_via_stdout_set_empty_json(
+        self, caplog: pytest.LogCaptureFixture, step_script_definition: StepScript_2023_09
+    ) -> None:
+        # Test that when an environment defines variables directly and as empty json it is processed properly
+
+        # GIVEN
+        environment = Environment_2023_09(
+            name="Env",
+            script=EnvironmentScript_2023_09(
+                actions=EnvironmentActions_2023_09(
+                    onEnter=Action_2023_09(
+                        command=sys.executable, args=["-c", "print('openjd_env: \"FOO=\"')"]
+                    )
+                )
+            ),
+        )
+        session_id = uuid.uuid4().hex
+        job_params = dict[str, ParameterValue]()
+        with Session(session_id=session_id, job_parameter_values=job_params) as session:
+            session.enter_environment(environment=environment)
+            while session.state == SessionState.RUNNING:
+                time.sleep(0.1)
+
+            # WHEN
+            session.run_task(
+                step_script=step_script_definition,
+                task_parameter_values=dict[str, ParameterValue](),
+            )
+            while session.state == SessionState.RUNNING:
+                time.sleep(0.1)
+
+            # THEN
+            assert "FOO=" in caplog.messages
+
+    @pytest.mark.usefixtures("caplog")  # builtin fixture
+    def test_def_via_stdout_set_empty(
+        self, caplog: pytest.LogCaptureFixture, step_script_definition: StepScript_2023_09
+    ) -> None:
+        # Test that when an environment defines variables directly and as empty string it is processed properly
+
+        # GIVEN
+        environment = Environment_2023_09(
+            name="Env",
+            script=EnvironmentScript_2023_09(
+                actions=EnvironmentActions_2023_09(
+                    onEnter=Action_2023_09(
+                        command=sys.executable, args=["-c", "print('openjd_env: FOO=')"]
+                    )
+                )
+            ),
+        )
+        session_id = uuid.uuid4().hex
+        job_params = dict[str, ParameterValue]()
+        with Session(session_id=session_id, job_parameter_values=job_params) as session:
+            session.enter_environment(environment=environment)
+            while session.state == SessionState.RUNNING:
+                time.sleep(0.1)
+
+            # WHEN
+            session.run_task(
+                step_script=step_script_definition,
+                task_parameter_values=dict[str, ParameterValue](),
+            )
+            while session.state == SessionState.RUNNING:
+                time.sleep(0.1)
+
+            # THEN
+            assert "FOO=" in caplog.messages
+
     @pytest.mark.usefixtures("caplog")  # builtin fixture
     def test_def_via_stdout_fails_session_action_on_error(
         self, caplog: pytest.LogCaptureFixture, step_script_definition: StepScript_2023_09