Restructure code to ensure coverage

Author: stxue1

cwltool/command_line_tool.py

@@ -888,26 +888,15 @@ def remove_prefix(s: str, prefix: str) -> str:
                         keydict[cls] = r
 
             # If there are environmental variables to preserve, add it to the key
-            env_def = cast(CWLObjectType, keydict.get("EnvVarRequirement", {}))
+            env_def = dict(cast(Mapping[str, str], keydict.get("EnvVarRequirement", {})))
             env_requirement, _ = self.get_requirement("EnvVarRequirement")
             if env_requirement:
                 for req in cast(list[CWLObjectType], env_requirement["envDef"]):
                     env_name = cast(str, req["envName"])
                     env_value = cast(str, req["envValue"])
                     env_def[env_name] = env_value
             if runtimeContext.preserve_environment is not None:
-                # This could either be a list or a dictionary
-                if isinstance(runtimeContext.preserve_environment, dict):
-                    for name, value in runtimeContext.preserve_environment.items():
-                        env_def[name] = value
-                elif isinstance(runtimeContext.preserve_environment, list):
-                    for key in runtimeContext.preserve_environment:
-                        try:
-                            env_def[key] = os.environ[key]
-                        except KeyError:
-                            _logger.warning(
-                                f"Attempting to preserve environment variable {key!r} which is not present"
-                            )
+                env_def = JobBase.extract_environment(runtimeContext, env_def)
 
             if env_def:
                 keydict["EnvVarRequirement"] = env_def

cwltool/job.py

@@ -469,6 +469,31 @@ def _preserve_environment_on_containers_warning(
         # By default, don't do anything; ContainerCommandLineJob below
         # will issue a warning.
 
+    @staticmethod
+    def extract_environment(
+        runtimeContext: RuntimeContext, envVarReq: Mapping[str, str]
+    ) -> dict[str, str]:
+        """
+        Extract environment variables that should be preserved
+        """
+        # Start empty
+        env: dict[str, str] = {}
+        # Preserve any env vars
+        if runtimeContext.preserve_entire_environment:
+            env.update(os.environ)
+        elif runtimeContext.preserve_environment:
+            for key in runtimeContext.preserve_environment:
+                try:
+                    env[key] = os.environ[key]
+                except KeyError:
+                    _logger.warning(
+                        f"Attempting to preserve environment variable {key!r} which is not present"
+                    )
+        # Apply EnvVarRequirement
+        env.update()
+
+        return env
+
     def prepare_environment(
         self, runtimeContext: RuntimeContext, envVarReq: Mapping[str, str]
     ) -> None:
@@ -481,26 +506,15 @@ def prepare_environment(
         """
         # Start empty
         env: dict[str, str] = {}
-
-        # Preserve any env vars
         if runtimeContext.preserve_entire_environment:
             self._preserve_environment_on_containers_warning()
-            env.update(os.environ)
         elif runtimeContext.preserve_environment:
             self._preserve_environment_on_containers_warning(runtimeContext.preserve_environment)
-            for key in runtimeContext.preserve_environment:
-                try:
-                    env[key] = os.environ[key]
-                except KeyError:
-                    _logger.warning(
-                        f"Attempting to preserve environment variable {key!r} which is not present"
-                    )
 
         # Set required env vars
         env.update(self._required_env())
 
-        # Apply EnvVarRequirement
-        env.update(envVarReq)
+        env.update(self.extract_environment(runtimeContext, envVarReq))
 
         # Set on ourselves
         self.environment = env

tests/test_examples.py

@@ -1346,6 +1346,7 @@ def test_cache_default_literal_file(tmp_path: Path, factor: str) -> None:
     assert "completed success" in stderr
     assert error_code == 0
 
+
 @pytest.mark.parametrize("factor", test_factors)
 def test_cache_environment_variable(tmp_path: Path, factor: str) -> None:
     """Ensure that changing the environment variables will result in different cache keys"""
@@ -1360,7 +1361,7 @@ def test_cache_environment_variable(tmp_path: Path, factor: str) -> None:
             "--outdir",
             str(tmp_path / "outdir"),
             get_data("tests/" + test_file),
-            get_data(f"tests/{test_job_file}")
+            get_data(f"tests/{test_job_file}"),
         ]
     )
     error_code, _, stderr = get_main_output(commands)
@@ -1382,7 +1383,7 @@ def test_cache_environment_variable(tmp_path: Path, factor: str) -> None:
             "--outdir",
             str(tmp_path / "outdir"),
             get_data("tests/" + test_file),
-            get_data(f"tests/{test_job_file}")
+            get_data(f"tests/{test_job_file}"),
         ]
     )
 
@@ -1392,6 +1393,7 @@ def test_cache_environment_variable(tmp_path: Path, factor: str) -> None:
     assert "Output of job will be cached in" in stderr
     assert error_code == 0, stderr
 
+
 def test_write_summary(tmp_path: Path) -> None:
     """Test --write-summary."""
     commands = [