From 786aead221a47b5c4f614676310406f2e01de2c1 Mon Sep 17 00:00:00 2001
From: Tony Meyer <tony.meyer@canonical.com>
Date: Wed, 20 Nov 2024 12:56:06 +1300
Subject: [PATCH] Provide a dict to _JujuContext directly.

No point manipulating os.environ - just provide the expected environment as a dictionary to the _JujuContext object.
---
 testing/src/scenario/ops_main_mock.py | 16 ++++++++++------
 testing/src/scenario/runtime.py       | 19 ++++---------------
 2 files changed, 14 insertions(+), 21 deletions(-)

diff --git a/testing/src/scenario/ops_main_mock.py b/testing/src/scenario/ops_main_mock.py
index 71b22e097..8b051b5a4 100644
--- a/testing/src/scenario/ops_main_mock.py
+++ b/testing/src/scenario/ops_main_mock.py
@@ -88,17 +88,16 @@ def setup_framework(
 ):
     from .mocking import _MockModelBackend
 
+    if juju_context is None:
+        juju_context = ops.jujucontext._JujuContext.from_dict(os.environ)
     model_backend = _MockModelBackend(
         state=state,
         event=event,
         context=context,
         charm_spec=charm_spec,
-        juju_context=ops.jujucontext._JujuContext.from_dict(os.environ)
-        if juju_context is None
-        else juju_context,
+        juju_context=juju_context,
     )
-    debug = "JUJU_DEBUG" in os.environ
-    setup_root_logging(model_backend, debug=debug)
+    setup_root_logging(model_backend, debug=juju_context.debug)
     # ops sets sys.excepthook to go to Juju's debug-log, but that's not useful
     # in a testing context, so reset it.
     sys.excepthook = sys.__excepthook__
@@ -190,12 +189,17 @@ def __init__(
         event: "_Event",
         context: "Context",
         charm_spec: "_CharmSpec[CharmType]",
+        juju_context: Optional[ops.jujucontext._JujuContext] = None,
     ):
         self.state = state
         self.event = event
         self.context = context
         self.charm_spec = charm_spec
-        self.juju_context = ops.jujucontext._JujuContext.from_dict(os.environ)
+        self.juju_context = (
+            ops.jujucontext._JujuContext.from_dict(os.environ)
+            if juju_context is None
+            else juju_context
+        )
 
         # set by setup()
         self.dispatcher: Optional[_Dispatcher] = None
diff --git a/testing/src/scenario/runtime.py b/testing/src/scenario/runtime.py
index 1206c119c..63a7d937a 100644
--- a/testing/src/scenario/runtime.py
+++ b/testing/src/scenario/runtime.py
@@ -5,7 +5,6 @@
 import copy
 import dataclasses
 import marshal
-import os
 import re
 import tempfile
 import typing
@@ -35,6 +34,7 @@
     NoTypeError,
     PreCommitEvent,
 )
+from ops.jujucontext import _JujuContext
 from ops.storage import NoSnapshotError, SQLiteStorage
 from ops.framework import _event_regex
 from ops._private.harness import ActionFailed
@@ -183,15 +183,6 @@ def __init__(
         self._app_name = app_name
         self._unit_id = unit_id
 
-    @staticmethod
-    def _cleanup_env(env: Dict[str, str]):
-        # TODO consider cleaning up env on __delete__, but ideally you should be
-        #  running this in a clean env or a container anyway.
-        # cleanup the env, in case we'll be firing multiple events, we don't want to pollute it.
-        for key in env:
-            # os.unsetenv does not always seem to work !?
-            del os.environ[key]
-
     def _get_event_env(self, state: "State", event: "_Event", charm_root: Path):
         """Build the simulated environment the operator framework expects."""
         env = {
@@ -430,7 +421,7 @@ def exec(
         Returns the 'output state', that is, the state as mutated by the charm during the
         event handling.
 
-        This will set the environment up and call ops.main.main().
+        This will set the environment up and call ops.main().
         After that it's up to ops.
         """
         # todo consider forking out a real subprocess and do the mocking by
@@ -457,7 +448,7 @@ def exec(
                 event=event,
                 charm_root=temporary_charm_root,
             )
-            os.environ.update(env)
+            juju_context = _JujuContext.from_dict(env)
 
             logger.info(" - Entering ops.main (mocked).")
             from .ops_main_mock import Ops  # noqa: F811
@@ -471,6 +462,7 @@ def exec(
                         self._charm_spec,
                         charm_type=self._wrap(charm_type),
                     ),
+                    juju_context=juju_context,
                 )
                 ops.setup()
 
@@ -489,9 +481,6 @@ def exec(
             finally:
                 logger.info(" - Exited ops.main.")
 
-                logger.info(" - Clearing env")
-                self._cleanup_env(env)
-
             logger.info(" - closing storage")
             output_state = self._close_storage(output_state, temporary_charm_root)