Copy the mocked Juju envvars into os.environ again.

Author: tonyandrewmeyer

testing/src/scenario/_runtime.py

@@ -322,10 +322,10 @@ def exec(
             juju_context = _JujuContext.from_dict(env)
             # We need to set JUJU_VERSION in os.environ, because charms may use
             # `JujuVersion.from_environ()` to get the (simulated) Juju version.
-            # All of the other environment variables are exposed to the charm
-            # through the framework machinery, so don't require this.
-            previous_juju_version = os.environ.get("JUJU_VERSION")
-            os.environ["JUJU_VERSION"] = env["JUJU_VERSION"]
+            # For consistency, we put all the other ones there too, although we'd
+            # like to change this in the future.
+            previous_env = os.environ.copy()
+            os.environ.update(env)
 
             logger.info(" - entering ops.main (mocked)")
             from ._ops_main_mock import Ops  # noqa: F811
@@ -352,10 +352,14 @@ def exec(
                 ) from e
 
             finally:
-                if previous_juju_version is None:
-                    del os.environ["JUJU_VERSION"]
-                else:
-                    os.environ["JUJU_VERSION"] = previous_juju_version
+                for key, value in os.environ.copy().items():
+                    if key in previous_env:
+                        os.environ[key] = value
+                    else:
+                        del os.environ[key]
+                for key, value in previous_env.items():
+                    if key not in os.environ:
+                        os.environ[key] = value
                 logger.info(" - exited ops.main")
 
         context.emitted_events.extend(captured)

testing/tests/test_runtime.py

@@ -110,7 +110,7 @@ def test_env_clean_on_charm_error():
             context=Context(my_charm_type, meta=meta),
         ) as manager:
             assert manager._juju_context.remote_app_name == remote_name
-            assert "JUJU_REMOTE_APP" not in os.environ
+            assert "JUJU_REMOTE_APP" in os.environ
             _ = 1 / 0  # raise some error
     # Ensure that some other error didn't occur (like AssertionError!).
     assert "ZeroDivisionError" in str(exc.value)