refactor!: use ops._main._Manager in Scenario

ops/_main.py

@@ -112,28 +112,6 @@ def _setup_event_links(charm_dir: Path, charm: 'ops.charm.CharmBase', juju_conte
             _create_event_link(charm, bound_event, link_to)
 
 
-def _emit_charm_event(charm: 'ops.charm.CharmBase', event_name: str, juju_context: _JujuContext):
-    """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.
-        juju_context: An instance of the _JujuContext class.
-    """
-    event_to_emit = None
-    try:
-        event_to_emit = getattr(charm.on, event_name)
-    except AttributeError:
-        logger.debug('Event %s not defined for %s.', event_name, charm)
-
-    # If the event is not supported by the charm implementation, do
-    # not error out or try to emit it. This is to support rollbacks.
-    if event_to_emit is not None:
-        args, kwargs = _get_event_args(charm, event_to_emit, juju_context)
-        logger.debug('Emitting Juju event %s.', event_name)
-        event_to_emit.emit(*args, **kwargs)
-
-
 def _get_event_args(
     charm: 'ops.charm.CharmBase',
     bound_event: 'ops.framework.BoundEvent',
@@ -401,8 +379,11 @@ def __init__(
         model_backend: Optional[ops.model._ModelBackend] = None,
         use_juju_for_storage: Optional[bool] = None,
         charm_state_path: str = CHARM_STATE_FILE,
+        juju_context: Optional[_JujuContext] = None,
     ):
-        self._juju_context = _JujuContext.from_dict(os.environ)
+        if juju_context is None:
+            juju_context = _JujuContext.from_dict(os.environ)
+        self._juju_context = juju_context
         self._charm_state_path = charm_state_path
         self._charm_class = charm_class
         if model_backend is None:
@@ -413,7 +394,7 @@ def __init__(
         self._setup_root_logging()
 
         self._charm_root = self._juju_context.charm_dir
-        self._charm_meta = CharmMeta.from_charm_root(self._charm_root)
+        self._charm_meta = self._load_charm_meta()
         self._use_juju_for_storage = use_juju_for_storage
 
         # Set up dispatcher, framework and charm objects.
@@ -423,6 +404,9 @@ def __init__(
         self.framework = self._make_framework(self.dispatcher)
         self.charm = self._make_charm(self.framework, self.dispatcher)
 
+    def _load_charm_meta(self):
+        return CharmMeta.from_charm_root(self._charm_root)
+
     def _make_charm(self, framework: 'ops.framework.Framework', dispatcher: _Dispatcher):
         charm = self._charm_class(framework)
         dispatcher.ensure_event_links(charm)
@@ -482,7 +466,7 @@ def _make_framework(self, dispatcher: _Dispatcher):
         # If we are in a RelationBroken event, we want to know which relation is
         # broken within the model, not only in the event's `.relation` attribute.
 
-        if self._juju_context.dispatch_path.endswith('-relation-broken'):
+        if self._juju_context.dispatch_path.endswith(('-relation-broken', '_relation_broken')):
             broken_relation_id = self._juju_context.relation_id
         else:
             broken_relation_id = None
@@ -515,19 +499,50 @@ def _emit(self):
             self.framework.reemit()
 
         # Emit the Juju event.
-        _emit_charm_event(self.charm, self.dispatcher.event_name, self._juju_context)
+        self._emit_charm_event(self.dispatcher.event_name)
         # Emit collect-status events.
         ops.charm._evaluate_status(self.charm)
 
+    def _get_event_to_emit(self, event_name: str) -> Optional[ops.framework.BoundEvent]:
+        try:
+            return getattr(self.charm.on, event_name)
+        except AttributeError:
+            logger.debug('Event %s not defined for %s.', event_name, self.charm)
+        return None
+
+    def _emit_charm_event(self, event_name: str):
+        """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.
+            juju_context: An instance of the _JujuContext class.
+        """
+        event_to_emit = self._get_event_to_emit(event_name)
+
+        # If the event is not supported by the charm implementation, do
+        # not error out or try to emit it. This is to support rollbacks.
+        if event_to_emit is None:
+            return
+
+        args, kwargs = _get_event_args(self.charm, event_to_emit, self._juju_context)
+        logger.debug('Emitting Juju event %s.', event_name)
+        event_to_emit.emit(*args, **kwargs)
+
     def _commit(self):
         """Commit the framework and gracefully teardown."""
         self.framework.commit()
 
+    def _close(self):
+        """Perform any necessary cleanup before the framework is closed."""
+        # Provided for child classes - nothing needs to be done in the base.
+
     def run(self):
         """Emit and then commit the framework."""
         try:
             self._emit()
             self._commit()
+            self._close()
         finally:
             self.framework.close()
 

ops/testing.py

@@ -177,8 +177,8 @@
         # monkeypatch it in, so that the ops.testing.ActionFailed object is the
         # one that we expect, even if people are mixing Harness and Scenario.
         # https://github.com/canonical/ops-scenario/issues/201
+        import scenario._runtime as _runtime
         import scenario.context as _context
-        import scenario.runtime as _runtime
 
         _context.ActionFailed = ActionFailed  # type: ignore[reportPrivateImportUsage]
         _runtime.ActionFailed = ActionFailed  # type: ignore[reportPrivateImportUsage]

test/test_main.py

@@ -97,7 +97,7 @@ def __init__(
 
 
 @patch('ops._main.setup_root_logging', new=lambda *a, **kw: None)  # type: ignore
-@patch('ops._main._emit_charm_event', new=lambda *a, **kw: None)  # type: ignore
+@patch('ops._main._Manager._emit_charm_event', new=lambda *a, **kw: None)  # type: ignore
 @patch('ops.charm._evaluate_status', new=lambda *a, **kw: None)  # type: ignore
 class TestCharmInit:
     @patch('sys.stderr', new_callable=io.StringIO)
@@ -235,11 +235,11 @@ def __init__(self, framework: ops.Framework):
                 dispatch.chmod(0o755)
 
             with patch.dict(os.environ, fake_environ):
-                with patch('ops._main._emit_charm_event') as mock_charm_event:
+                with patch('ops._main._Manager._emit_charm_event') as mock_charm_event:
                     ops.main(MyCharm)
 
         assert mock_charm_event.call_count == 1
-        return mock_charm_event.call_args[0][1]
+        return mock_charm_event.call_args[0][0]
 
     def test_most_legacy(self):
         """Without dispatch, sys.argv[0] is used."""

test/test_main_invocation.py

@@ -24,7 +24,7 @@
 @pytest.fixture
 def charm_env(monkeypatch: pytest.MonkeyPatch, tmp_path: Path):
     monkeypatch.setattr('sys.argv', ('hooks/install',))
-    monkeypatch.setattr('ops._main._emit_charm_event', Mock())
+    monkeypatch.setattr('ops._main._Manager._emit_charm_event', Mock())
     monkeypatch.setattr('ops._main._Manager._setup_root_logging', Mock())
     monkeypatch.setattr('ops.charm._evaluate_status', Mock())
     monkeypatch.setenv('JUJU_CHARM_DIR', str(tmp_path))

testing/src/scenario/_consistency_checker.py

@@ -39,7 +39,7 @@
 )
 
 from .errors import InconsistentScenarioError
-from .runtime import logger as scenario_logger
+from ._runtime import logger as scenario_logger
 from .state import (
     CharmType,
     PeerRelation,
@@ -179,6 +179,11 @@ def check_event_consistency(
         # skip everything here. Perhaps in the future, custom events could
         # optionally include some sort of state metadata that made testing
         # consistency possible?
+        warnings.append(
+            "this is a custom event; if its name makes it look like a builtin one "
+            "(for example, a relation event, or a workload event), you might get some false-negative "
+            "consistency checks.",
+        )
         return Results(errors, warnings)
 
     if event._is_relation_event:

testing/src/scenario/_ops_main_mock.py

@@ -0,0 +1,197 @@
+#!/usr/bin/env python3
+# Copyright 2023 Canonical Ltd.
+# See LICENSE file for licensing details.
+
+import dataclasses
+import marshal
+import re
+import sys
+from typing import TYPE_CHECKING, Any, Dict, FrozenSet, List, Sequence, Set
+
+import ops
+import ops.jujucontext
+import ops.storage
+
+from ops.framework import _event_regex
+from ops._main import _Dispatcher, _Manager
+from ops._main import logger as ops_logger
+
+from .errors import BadOwnerPath, NoObserverError
+from .logger import logger as scenario_logger
+from .mocking import _MockModelBackend
+from .state import CharmType, StoredState, DeferredEvent
+
+if TYPE_CHECKING:  # pragma: no cover
+    from .context import Context
+    from .state import State, _CharmSpec, _Event
+
+EVENT_REGEX = re.compile(_event_regex)
+STORED_STATE_REGEX = re.compile(
+    r"((?P<owner_path>.*)\/)?(?P<_data_type_name>\D+)\[(?P<name>.*)\]",
+)
+
+logger = scenario_logger.getChild("ops_main_mock")
+
+# pyright: reportPrivateUsage=false
+
+
+class UnitStateDB:
+    """Wraps the unit-state database with convenience methods for adjusting the state."""
+
+    def __init__(self, underlying_store: ops.storage.SQLiteStorage):
+        self._db = underlying_store
+
+    def get_stored_states(self) -> FrozenSet["StoredState"]:
+        """Load any StoredState data structures from the db."""
+        db = self._db
+        stored_states: Set[StoredState] = set()
+        for handle_path in db.list_snapshots():
+            if not EVENT_REGEX.match(handle_path) and (
+                match := STORED_STATE_REGEX.match(handle_path)
+            ):
+                stored_state_snapshot = db.load_snapshot(handle_path)
+                kwargs = match.groupdict()
+                sst = StoredState(content=stored_state_snapshot, **kwargs)
+                stored_states.add(sst)
+
+        return frozenset(stored_states)
+
+    def get_deferred_events(self) -> List["DeferredEvent"]:
+        """Load any DeferredEvent data structures from the db."""
+        db = self._db
+        deferred: List[DeferredEvent] = []
+        for handle_path in db.list_snapshots():
+            if EVENT_REGEX.match(handle_path):
+                notices = db.notices(handle_path)
+                for handle, owner, observer in notices:
+                    try:
+                        snapshot_data = db.load_snapshot(handle)
+                    except ops.storage.NoSnapshotError:
+                        snapshot_data: Dict[str, Any] = {}
+
+                    event = DeferredEvent(
+                        handle_path=handle,
+                        owner=owner,
+                        observer=observer,
+                        snapshot_data=snapshot_data,
+                    )
+                    deferred.append(event)
+
+        return deferred
+
+    def apply_state(self, state: "State"):
+        """Add DeferredEvent and StoredState from this State instance to the storage."""
+        db = self._db
+        for event in state.deferred:
+            db.save_notice(event.handle_path, event.owner, event.observer)
+            try:
+                marshal.dumps(event.snapshot_data)
+            except ValueError as e:
+                raise ValueError(
+                    f"unable to save the data for {event}, it must contain only simple types.",
+                ) from e
+            db.save_snapshot(event.handle_path, event.snapshot_data)
+
+        for stored_state in state.stored_states:
+            db.save_snapshot(stored_state._handle_path, stored_state.content)
+
+
+class Ops(_Manager):
+    """Class to manage stepping through ops setup, event emission and framework commit."""
+
+    def __init__(
+        self,
+        state: "State",
+        event: "_Event",
+        context: "Context[CharmType]",
+        charm_spec: "_CharmSpec[CharmType]",
+        juju_context: ops.jujucontext._JujuContext,
+    ):
+        self.state = state
+        self.event = event
+        self.context = context
+        self.charm_spec = charm_spec
+        self.store = None
+
+        model_backend = _MockModelBackend(
+            state=state,
+            event=event,
+            context=context,
+            charm_spec=charm_spec,
+            juju_context=juju_context,
+        )
+
+        super().__init__(
+            self.charm_spec.charm_type, model_backend, juju_context=juju_context
+        )
+
+    def _load_charm_meta(self):
+        metadata = (self._charm_root / "metadata.yaml").read_text()
+        actions_meta = self._charm_root / "actions.yaml"
+        if actions_meta.exists():
+            actions_metadata = actions_meta.read_text()
+        else:
+            actions_metadata = None
+
+        return ops.CharmMeta.from_yaml(metadata, actions_metadata)
+
+    def _setup_root_logging(self):
+        # Ops sets sys.excepthook to go to Juju's debug-log, but that's not
+        # useful in a testing context, so we reset it here.
+        super()._setup_root_logging()
+        sys.excepthook = sys.__excepthook__
+
+    def _make_storage(self, _: _Dispatcher):
+        # TODO: add use_juju_for_storage support
+        # TODO: Pass a charm_state_path that is ':memory:' when appropriate.
+        charm_state_path = self._charm_root / self._charm_state_path
+        storage = ops.storage.SQLiteStorage(charm_state_path)
+        logger.info("Copying input state to storage.")
+        self.store = UnitStateDB(storage)
+        self.store.apply_state(self.state)
+        return storage
+
+    def _get_event_to_emit(self, event_name: str):
+        owner = (
+            self._get_owner(self.charm, self.event.owner_path)
+            if self.event
+            else self.charm.on
+        )
+
+        try:
+            event_to_emit = getattr(owner, event_name)
+        except AttributeError:
+            ops_logger.debug("Event %s not defined for %s.", event_name, self.charm)
+            raise NoObserverError(
+                f"Cannot fire {event_name!r} on {owner}: "
+                f"invalid event (not on charm.on).",
+            )
+        return event_to_emit
+
+    @staticmethod
+    def _get_owner(root: Any, path: Sequence[str]) -> ops.ObjectEvents:
+        """Walk path on root to an ObjectEvents instance."""
+        obj = root
+        for step in path:
+            try:
+                obj = getattr(obj, step)
+            except AttributeError:
+                raise BadOwnerPath(
+                    f"event_owner_path {path!r} invalid: {step!r} leads to nowhere.",
+                )
+        if not isinstance(obj, ops.ObjectEvents):
+            raise BadOwnerPath(
+                f"event_owner_path {path!r} invalid: does not lead to "
+                f"an ObjectEvents instance.",
+            )
+        return obj
+
+    def _close(self):
+        """Now that we're done processing this event, read the charm state and expose it."""
+        logger.info("Copying storage to output state.")
+        assert self.store is not None
+        deferred = self.store.get_deferred_events()
+        stored_state = self.store.get_stored_states()
+        self.state = dataclasses.replace(
+            self.state, deferred=deferred, stored_states=stored_state
+        )