Address PR feedback: rollback partial env, fail on bad templates

- TmpEnv.__enter__ rolls back already-set env vars if swap_env raises
  mid-loop, preventing leaked env state on enter failures
- Remove silent TemplateError swallow: malformed templates now raise
  LookupError with the source string and error detail instead of
  silently passing through as literals
- Convert tests to plain pytest classes (drop unittest.TestCase)
- Add test_tmpenv_rollback_on_error for partial modification safety

Author: anticomputer

src/seclab_taskflow_agent/env_utils.py

@@ -52,11 +52,9 @@ def swap_env(s: str, context: dict[str, Any] | None = None) -> str:
         }
         return template.render(**render_context)
     except jinja2.UndefinedError as e:
-        # Convert Jinja undefined to LookupError for compatibility
         raise LookupError(str(e))
-    except jinja2.TemplateError:
-        # Not a template or failed to render, return as-is
-        return s
+    except jinja2.TemplateError as e:
+        raise LookupError(f"Template rendering failed for: {s!r}: {e}")
 
 
 class TmpEnv:
@@ -69,8 +67,18 @@ def __init__(self, env: dict[str, str],
         self.restore_env = dict(os.environ)
 
     def __enter__(self) -> None:
-        for k, v in self.env.items():
-            os.environ[k] = swap_env(v, self.context)
+        applied: list[str] = []
+        try:
+            for k, v in self.env.items():
+                os.environ[k] = swap_env(v, self.context)
+                applied.append(k)
+        except Exception:
+            for k in applied:
+                if k in self.restore_env:
+                    os.environ[k] = self.restore_env[k]
+                else:
+                    os.environ.pop(k, None)
+            raise
 
     def __exit__(self, exc_type: type | None, exc_val: BaseException | None, exc_tb: Any | None) -> None:
         for k, v in self.env.items():

tests/test_env_utils.py

@@ -4,14 +4,13 @@
 """Tests for env_utils: swap_env and TmpEnv with globals context."""
 
 import os
-import unittest
 
 import pytest
 
 from seclab_taskflow_agent.env_utils import TmpEnv, swap_env
 
 
-class TestSwapEnv(unittest.TestCase):
+class TestSwapEnv:
     """Tests for swap_env template rendering."""
 
     def test_plain_string_unchanged(self):
@@ -51,7 +50,7 @@ def test_no_context_backward_compat(self):
         assert swap_env("plain") == "plain"
 
 
-class TestTmpEnv(unittest.TestCase):
+class TestTmpEnv:
     """Tests for TmpEnv context manager with globals."""
 
     def test_globals_rendered_in_env_block(self):
@@ -77,3 +76,12 @@ def test_tmpenv_restores_original(self):
             assert os.environ["RESTORE_TEST"] == "overwritten"
         assert os.environ["RESTORE_TEST"] == "original"
         del os.environ["RESTORE_TEST"]
+
+    def test_tmpenv_rollback_on_error(self):
+        """Partial env modification is rolled back if swap_env raises."""
+        env = {"GOOD_KEY": "value", "BAD_KEY": "{{ globals.missing }}"}
+        with pytest.raises(LookupError):
+            with TmpEnv(env):
+                pass
+        assert "GOOD_KEY" not in os.environ
+        assert "BAD_KEY" not in os.environ