refactor: use _JujuContext in Scenario (#1459)

Properly use the `ops.jujucontext._JujuContext` class in Scenario.

When `_JujuContext` was introduced in ops, we added some quick
compatibility code to Scenario that made use of it, but it was basically
just recreating the context object from the environment everywhere.

This PR tidies things up in two main ways:
* Create a single `_JujuContext` object and pass it around where it's
required. This also sets the stage for using the ops `_Manager` class in
a later PR.
* Rather than updating and then later (imprecisely) reverting
`os.environ`, just pass the mock environment directly when creating the
`_JujuContext` object.

Note that to maintain backwards compatibility of public methods the
argument is optional in some places and will fall back to the
environment - this won't be exactly the same, since we aren't
manipulating the environment any more, but the signatures are the same.
I don't think it's worth jumping through the hoops to have it entirely
the same behaviour (by having the fallback do environment patching). It
seems unlikely anyone is actually using these, even though they are
public (the only candidate I can think of is Catan, but that's pinned to
6.x). Perhaps, given that the *behaviour* ends up potentially changing
(depending on how the caller is using things), it's not worth keeping
the signatures the same? I'm interested in reviewer's opinions on this.

---------

Co-authored-by: Ben Hoyt <[email protected]>

testing/src/scenario/mocking.py

@@ -33,6 +33,7 @@
     ModelError,
 )
 from ops._private.harness import ExecArgs, _TestingPebbleClient
+from ops.jujucontext import _JujuContext
 from ops.model import CloudSpec as CloudSpec_Ops
 from ops.model import Port as Port_Ops
 from ops.model import Secret as Secret_Ops  # lol
@@ -130,8 +131,9 @@ def __init__(
         event: "_Event",
         charm_spec: "_CharmSpec[CharmType]",
         context: "Context",
+        juju_context: "_JujuContext",
     ):
-        super().__init__()
+        super().__init__(juju_context=juju_context)
         self._state = state
         self._event = event
         self._context = context

testing/src/scenario/ops_main_mock.py

@@ -30,17 +30,6 @@
 # pyright: reportPrivateUsage=false
 
 
-# TODO: Use ops.jujucontext's _JujuContext.charm_dir.
-def _get_charm_dir():
-    charm_dir = os.environ.get("JUJU_CHARM_DIR")
-    if charm_dir is None:
-        # Assume $JUJU_CHARM_DIR/lib/op/main.py structure.
-        charm_dir = pathlib.Path(f"{__file__}/../../..").resolve()
-    else:
-        charm_dir = pathlib.Path(charm_dir).resolve()
-    return charm_dir
-
-
 def _get_owner(root: Any, path: Sequence[str]) -> ops.ObjectEvents:
     """Walk path on root to an ObjectEvents instance."""
     obj = root
@@ -62,14 +51,16 @@ def _get_owner(root: Any, path: Sequence[str]) -> ops.ObjectEvents:
 def _emit_charm_event(
     charm: "CharmBase",
     event_name: str,
+    juju_context: ops.jujucontext._JujuContext,
     event: Optional["_Event"] = None,
 ):
     """Emits a charm event based on a Juju event name.
 
     Args:
         charm: A charm instance to emit an event from.
         event_name: A Juju event name to emit on a charm.
-        event_owner_path: Event source lookup path.
+        event: Event to emit.
+        juju_context: Juju context to use for the event.
     """
     owner = _get_owner(charm, event.owner_path) if event else charm.on
 
@@ -82,11 +73,7 @@ def _emit_charm_event(
             f"invalid event (not on charm.on).",
         )
 
-    args, kwargs = _get_event_args(
-        charm,
-        event_to_emit,
-        ops.jujucontext._JujuContext.from_dict(os.environ),
-    )
+    args, kwargs = _get_event_args(charm, event_to_emit, juju_context)
     ops_logger.debug("Emitting Juju event %s.", event_name)
     event_to_emit.emit(*args, **kwargs)
 
@@ -97,17 +84,20 @@ def setup_framework(
     event: "_Event",
     context: "Context",
     charm_spec: "_CharmSpec[CharmType]",
+    juju_context: Optional[ops.jujucontext._JujuContext] = None,
 ):
     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=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__
@@ -172,18 +162,20 @@ def setup(
     event: "_Event",
     context: "Context",
     charm_spec: "_CharmSpec[CharmType]",
+    juju_context: Optional[ops.jujucontext._JujuContext] = None,
 ):
     """Setup dispatcher, framework and charm objects."""
     charm_class = charm_spec.charm_type
-    charm_dir = _get_charm_dir()
+    if juju_context is None:
+        juju_context = ops.jujucontext._JujuContext.from_dict(os.environ)
+    charm_dir = juju_context.charm_dir
 
-    dispatcher = _Dispatcher(
-        charm_dir,
-        ops.jujucontext._JujuContext.from_dict(os.environ),
-    )
+    dispatcher = _Dispatcher(charm_dir, juju_context)
     dispatcher.run_any_legacy_hook()
 
-    framework = setup_framework(charm_dir, state, event, context, charm_spec)
+    framework = setup_framework(
+        charm_dir, state, event, context, charm_spec, juju_context
+    )
     charm = setup_charm(charm_class, framework, dispatcher)
     return dispatcher, framework, charm
 
@@ -197,11 +189,15 @@ 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
+        if juju_context is None:
+            juju_context = ops.jujucontext._JujuContext.from_dict(os.environ)
+        self.juju_context = juju_context
 
         # set by setup()
         self.dispatcher: Optional[_Dispatcher] = None
@@ -220,6 +216,7 @@ def setup(self):
             self.event,
             self.context,
             self.charm_spec,
+            self.juju_context,
         )
 
     def emit(self):
@@ -236,7 +233,9 @@ def emit(self):
             if not dispatcher.is_restricted_context():
                 framework.reemit()
 
-            _emit_charm_event(charm, dispatcher.event_name, self.event)
+            _emit_charm_event(
+                charm, dispatcher.event_name, self.juju_context, self.event
+            )
 
         except Exception:
             framework.close()

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)
 

testing/tests/test_runtime.py

@@ -85,11 +85,11 @@ def test_unit_name(app_name, unit_id):
         state=State(),
         event=_Event("start"),
         context=Context(my_charm_type, meta=meta),
-    ) as ops:
-        assert ops.charm.unit.name == f"{app_name}/{unit_id}"
+    ) as manager:
+        assert manager.charm.unit.name == f"{app_name}/{unit_id}"
 
 
-def test_env_cleanup_on_charm_error():
+def test_env_clean_on_charm_error():
     meta = {"name": "frank", "requires": {"box": {"interface": "triangle"}}}
 
     my_charm_type = charm_type()
@@ -101,14 +101,19 @@ def test_env_cleanup_on_charm_error():
         ),
     )
 
-    rel = Relation("box")
-    with pytest.raises(UncaughtCharmError):
+    remote_name = "ava"
+    rel = Relation("box", remote_app_name=remote_name)
+    with pytest.raises(UncaughtCharmError) as exc:
         with runtime.exec(
             state=State(relations={rel}),
             event=_Event("box_relation_changed", relation=rel),
             context=Context(my_charm_type, meta=meta),
-        ):
-            assert os.getenv("JUJU_REMOTE_APP")
+        ) as manager:
+            assert manager.juju_context.remote_app_name == remote_name
+            assert "JUJU_REMOTE_APP" not in os.environ
             _ = 1 / 0  # raise some error
+    # Ensure that some other error didn't occur (like AssertionError!).
+    assert "ZeroDivisionError" in str(exc.value)
 
+    # Ensure that the Juju environment didn't leak into the outside one.
     assert os.getenv("JUJU_REMOTE_APP", None) is None