Repository: canonical/operator (python)

Issue:
Run deferred Juju events on a fresh charm instance
The main benefit here would be that it would be simpler to reason about (Juju) event behaviour, because it would be consistent both when an event was deferred and when it is executed normally.

This would not apply to custom events, which are executed synchronously, and are already different from Juju events. This would also not apply to `LifecycleEvents` (framework generated events), because those cannot be deferred.

We will need to use the same charm instance for multiple deferred (event, handler) pairs when more than one handler defers the event. If something like this happens:

* original Juju event: observer A defers, observer B does not
* next Juju event of same type: deferred observer A defers, observer A defers, observer B defers

I don't think we can tell that these were different Juju events, and we would maybe collapse the duplicates (?) and probably have to use the same charm instance for both observers. We would otherwise need to add some sort of ID to the deferred event to be able to group them correctly.

I believe the cleanest change here is to lift out the deferred `reemit` from `_Manager` in [ops/main.py], and handle the deferred events in `main()`, creating new `_Managers` for each one (we need to check that `setup_root_logging` can be called twice, and avoid calling legacy hooks twice - probably the legacy hook should be moved out of `_Manager` too). This provides a fresh `Model`, `Framework` and `_Dispatcher` as well, which is also most consistent (and avoids having to have the framework forget/unobserve the old charm).

There would be some small behaviour changes here: the pre-commit and commit events would get emitted twice (just as if the event had not been deferred, but not like now), which I think is fine, and we would collect the status twice (again just as if the event had not been deferred, but not like now), which might break some expectations.

See also the discussion in #736. Note that we will not be able to make Harness mimic this behaviour (because it would break too many existing charm tests) but we will be encouraging using Scenario, which can handle it (right now Scenario copies the ops behaviour, saving the deferred events to a queue that ops then runs ahead of the Juju event, but the Scenario API would not need changing in order to match the new ops behaviour).

Pull request description:
feat: run deferred events with fresh charm instances
Create a new charm object for each emission of a deferred event, and for the event that triggered the ops process.

When using the `testing.Context` as a context manager, the charm that is made available is the one for the primary event. If there is demand in the future, we could add a mechanism for accessing the other charm objects for manipulation prior to running the deferred events, or inspection after.

This will break charms that are relying on the charm object to be the same object for deferred and the primary event. However, we have always told charmers to not do that, and it's risky behaviour because it'll differ based on whether it's the non-deferred or deferred handler call, so we're ok with this change.

The other significant change for charms is that, because we need to commit the framework before finishing up with each charm, the `pre-commit` and `commit` Lifecycle events are emitted after each deferred event as well as after the primary event. I'm only aware of two charms that use these:
* ops-sunbeam, which uses `commit` to put statuses in stored state. Running more than once should be fine, and the performance impact not noticeable.
* `statustest` (of @benhoyt's test-charms collection), which seems almost identical to the sunbeam one. I assume this isn't really used, and it wouldn't break anyway.

There are three main parts of this implementation:
* Fairly minor adjustments to `_main`, particularly `_Manager`, so that the charm is passed around and the framework is accessed via the charm, rather than creating the framework and charm once, and the actual creation of each charm prior to re-emitting.
* Functionality to "forget" charms. The framework registers event types on the class, not the instance, so creating multiple charm instances, even on different `Framework`s, is not normally possible (Scenario gets around this by wrapping the event source in a new class for each run). `CharmBase` gains a `_destory_charm` function that undoes all of the event registration done in `__init__` (for all the dynamic events from relations, containers, and so forth). I considered using `__del__` for this, but it's awkward to cleanly remove everything when destruction is happening, and we don't really want charmers to start creating more than one charm instance, so this leaves the functionality private.
* Reworking the Scenario "capture events" mechanism, to use a `Framework` subclass, rather than a context manager. This needed changes to handle capturing deferred events, but it also fixes a bug. Previously, the `Framework` *class* was patched, with methods that used local variables to determine which events would be captured: this meant that the `capture_events()` context manager was not threadsafe.

Fixes #1174

Patch:
diff --git a/ops/_main.py b/ops/_main.py
index 975085f..fed072e 100644
--- a/ops/_main.py
+++ b/ops/_main.py
@@ -16,11 +16,11 @@
 
 import logging
 import os
+import pathlib
 import shutil
 import subprocess
 import sys
 import warnings
-from pathlib import Path
 from typing import Any, Dict, List, Optional, Tuple, Type, Union, cast
 
 from . import charm as _charm
@@ -37,7 +37,7 @@ CHARM_STATE_FILE = '.unit-state.db'
 logger = logging.getLogger()
 
 
-def _exe_path(path: Path) -> Optional[Path]:
+def _exe_path(path: pathlib.Path) -> Optional[pathlib.Path]:
     """Find and return the full path to the given binary.
 
     Here path is the absolute path to a binary, but might be missing an extension.
@@ -45,7 +45,7 @@ def _exe_path(path: Path) -> Optional[Path]:
     p = shutil.which(path.name, mode=os.F_OK, path=str(path.parent))
     if p is None:
         return None
-    return Path(p)
+    return pathlib.Path(p)
 
 
 def _get_event_args(
@@ -148,13 +148,13 @@ class _Dispatcher:
 
     event_name: str
 
-    def __init__(self, charm_dir: Path, juju_context: _JujuContext):
+    def __init__(self, charm_dir: pathlib.Path, juju_context: _JujuContext):
         self._juju_context = juju_context
         self._charm_dir = charm_dir
-        self._exec_path = Path(self._juju_context.dispatch_path or sys.argv[0])
+        self._exec_path = pathlib.Path(self._juju_context.dispatch_path or sys.argv[0])
 
         # Grab the correct hook from JUJU_DISPATCH_PATH, e.g. hooks/install.
-        self._dispatch_path = Path(self._juju_context.dispatch_path)
+        self._dispatch_path = pathlib.Path(self._juju_context.dispatch_path)
 
         if 'OPERATOR_DISPATCH' in os.environ:
             logger.debug('Charm called itself via %s.', self._dispatch_path)
@@ -179,7 +179,7 @@ class _Dispatcher:
             logger.warning('Legacy %s exists but is not executable.', self._dispatch_path)
             return
 
-        if dispatch_path.resolve() == Path(sys.argv[0]).resolve():
+        if dispatch_path.resolve() == pathlib.Path(sys.argv[0]).resolve():
             logger.debug('Legacy %s is just a link to ourselves.', self._dispatch_path)
             return
 
@@ -197,7 +197,7 @@ class _Dispatcher:
         else:
             logger.debug('Legacy %s exited with status 0.', self._dispatch_path)
 
-    def _set_name_from_path(self, path: Path):
+    def _set_name_from_path(self, path: pathlib.Path):
         """Sets the name attribute to that which can be inferred from the given path."""
         name = path.name.replace('-', '_')
         if path.parent.name == 'actions':
@@ -215,7 +215,7 @@ class _Dispatcher:
 
 
 def _should_use_controller_storage(
-    db_path: Path, meta: _charm.CharmMeta, juju_context: _JujuContext
+    db_path: pathlib.Path, meta: _charm.CharmMeta, juju_context: _JujuContext
 ) -> bool:
     """Figure out whether we want to use controller storage or not."""
     # if local state has been used previously, carry on using that
@@ -255,27 +255,35 @@ class _Manager:
       - the event that Juju is emitting on us
       - the charm instance (user-facing)
     - emit: core user-facing lifecycle step. Consists of:
-      - re-emit any deferred events found in the storage
       - emit the Juju event to the charm
         - emit any custom events emitted by the charm during this phase
       - emit  the ``collect-status`` events
     - commit: responsible for:
       - store any events deferred throughout this execution
       - graceful teardown of the storage
+
+    The above steps are first run for any deferred notices found in the storage
+    (all three steps for each notice, except for emitting status collection
+    events), and then run a final time for the Juju event that triggered this
+    execution.
     """
 
     def __init__(
         self,
         charm_class: Type['_charm.CharmBase'],
-        model_backend: Optional[_model._ModelBackend] = None,
+        juju_context: _JujuContext,
         use_juju_for_storage: Optional[bool] = None,
         charm_state_path: str = CHARM_STATE_FILE,
-        juju_context: Optional[_JujuContext] = None,
     ):
-        from . import tracing  # break circular import
+        # The context is shared across deferred events and the Juju event. Any
+        # data from the context that is event-specific must be included in the
+        # event object snapshot/restore rather than re-read from the context.
+        # Data not connected to the event (debug settings, the Juju version, the
+        # app and unit name, and so forth) will be the *current* data, not the
+        # data at the time the event was deferred -- this aligns with the data
+        # from hook tools.
 
-        if juju_context is None:
-            juju_context = _JujuContext.from_dict(os.environ)
+        from . import tracing  # break circular import
 
         try:
             name = charm_class.__name__
@@ -289,9 +297,6 @@ class _Manager:
         self._tracing_context.__enter__()
         self._charm_state_path = charm_state_path
         self._charm_class = charm_class
-        if model_backend is None:
-            model_backend = _model._ModelBackend(juju_context=self._juju_context)
-        self._model_backend = model_backend
 
         # Do this as early as possible to be sure to catch the most logs.
         self._setup_root_logging()
@@ -300,28 +305,57 @@ class _Manager:
         self._charm_meta = self._load_charm_meta()
         self._use_juju_for_storage = use_juju_for_storage
 
-        # Set up dispatcher, framework and charm objects.
-        self.dispatcher = _Dispatcher(self._charm_root, self._juju_context)
-        self.dispatcher.run_any_legacy_hook()
+        # Handle legacy hooks - this is only done once, not with each deferred
+        # event.
+        self._dispatcher = _Dispatcher(self._charm_root, self._juju_context)
+        self._dispatcher.run_any_legacy_hook()
+
+        # Storage is shared across all events, so we create it once here.
+        self._storage = self._make_storage()
+
+        self.run_deferred()
 
-        self.framework = self._make_framework(self.dispatcher)
-        self.charm = self._charm_class(self.framework)
+        # This is the charm for the Juju event. We create it here so that it's
+        # available for pre-emit adjustments when being used in testing.
+        self.charm = self._make_charm(self._dispatcher.event_name)
 
     def _load_charm_meta(self):
         return _charm.CharmMeta.from_charm_root(self._charm_root)
 
+    def _make_model_backend(self):
+        # model._ModelBackend is stateless and can be reused across events.
+        # However, in testing (both Harness and Scenario) the backend stores all
+        # the state that is normally in Juju. To be consistent, we create a new
+        # backend object even in production code.
+        return _model._ModelBackend(juju_context=self._juju_context)
+
+    def _make_charm(self, event_name: str):
+        framework = self._build_framework(event_name)
+        return self._charm_class(framework)
+
     def _setup_root_logging(self):
         # For actions, there is a communication channel with the user running the
         # action, so we want to send exception details through stderr, rather than
         # only to juju-log as normal.
         handling_action = self._juju_context.action_name is not None
+        # We don't really want to have a different backend here than when
+        # running the event. However, we need to create a new backend for each
+        # event and want the logging set up before we are ready to emit an
+        # event. In practice, this isn't a problem:
+        # * for model._ModelBackend, `juju_log` calls out directly to the hook
+        #   tool; it's effectively a staticmethod.
+        # * for _private.harness._TestingModelBackend, `juju_log` is not
+        #   implemented, and the logging is never configured.
+        # * for scenario.mocking._MockModelBackend, `juju_log` sends the logging
+        #   through to the `Context` object, which will be the same for all
+        #   events.
         setup_root_logging(
-            self._model_backend, debug=self._juju_context.debug, exc_stderr=handling_action
+            self._make_model_backend(), debug=self._juju_context.debug, exc_stderr=handling_action
         )
 
         logger.debug('ops %s up and running.', version)
 
-    def _make_storage(self, dispatcher: _Dispatcher):
+    def _make_storage(self):
         charm_state_path = self._charm_root / self._charm_state_path
 
         use_juju_for_storage = self._use_juju_for_storage
@@ -341,15 +375,12 @@ class _Manager:
                 category=DeprecationWarning,
             )
 
-        if use_juju_for_storage and dispatcher.is_restricted_context():
-            # TODO: jam 2020-06-30 This unconditionally avoids running a collect metrics event
-            #  Though we eventually expect that Juju will run collect-metrics in a
-            #  non-restricted context. Once we can determine that we are running
-            #  collect-metrics in a non-restricted context, we should fire the event as normal.
+        if use_juju_for_storage and self._dispatcher.is_restricted_context():
+            # collect-metrics is going away in Juju 4.0, and restricted context
+            # with it, so we don't need this to be particularly generic.
             logger.debug(
-                '"%s" is not supported when using Juju for storage\n'
+                '"collect_metrics" is not supported when using Juju for storage\n'
                 'see: https://github.com/canonical/operator/issues/348',
-                dispatcher.event_name,
             )
             # Note that we don't exit nonzero, because that would cause Juju to rerun the hook
             raise _Abort(0)
@@ -360,7 +391,12 @@ class _Manager:
             store = _storage.SQLiteStorage(charm_state_path)
         return store
 
-    def _make_framework(self, dispatcher: _Dispatcher):
+    def _make_framework(self, *args: Any, **kwargs: Any):
+        # A wrapper so that the testing subclasses can easily override which
+        # framework class is created.
+        return _framework.Framework(*args, **kwargs)
+
+    def _build_framework(self, event_name: str):
         # 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.
 
@@ -369,53 +405,45 @@ class _Manager:
         else:
             broken_relation_id = None
 
+        model_backend = self._make_model_backend()
         model = _model.Model(
-            self._charm_meta, self._model_backend, broken_relation_id=broken_relation_id
+            self._charm_meta, model_backend, broken_relation_id=broken_relation_id
         )
-        store = self._make_storage(dispatcher)
-        framework = _framework.Framework(
-            store,
+        framework = self._make_framework(
+            self._storage,
             self._charm_root,
             self._charm_meta,
             model,
-            event_name=dispatcher.event_name,
+            event_name=event_name,
             juju_debug_at=self._juju_context.debug_at,
         )
         framework.set_breakpointhook()
         return framework
 
-    def _emit(self):
+    def _emit(self, charm: _charm.CharmBase, event_name: str):
         """Emit the event on the charm."""
-        # TODO: Remove the collect_metrics check below as soon as the relevant
-        #       Juju changes are made. Also adjust the docstring on
-        #       EventBase.defer().
-        #
-        # Skip reemission of deferred events for collect-metrics events because
-        # they do not have the full access to all hook tools.
-        if not self.dispatcher.is_restricted_context():
-            # Re-emit any deferred events from the previous run.
-            self.framework.reemit()
-
         # Emit the Juju event.
-        self._emit_charm_event(self.dispatcher.event_name)
+        self._emit_charm_event(charm, event_name)
         # Emit collect-status events.
-        _charm._evaluate_status(self.charm)
+        _charm._evaluate_status(charm)
 
-    def _get_event_to_emit(self, event_name: str) -> Optional[_framework.BoundEvent]:
+    def _get_event_to_emit(
+        self, charm: _charm.CharmBase, event_name: str
+    ) -> Optional[_framework.BoundEvent]:
         try:
-            return getattr(self.charm.on, event_name)
+            return getattr(charm.on, event_name)
         except AttributeError:
-            logger.debug('Event %s not defined for %s.', event_name, self.charm)
+            logger.debug('Event %s not defined for %s.', event_name, charm)
         return None
 
     def _get_event_args(
-        self, bound_event: '_framework.BoundEvent'
+        self, charm: _charm.CharmBase, bound_event: '_framework.BoundEvent'
     ) -> Tuple[List[Any], Dict[str, Any]]:
         # A wrapper so that the testing subclasses can easily override the
         # behaviour.
-        return _get_event_args(self.charm, bound_event, self._juju_context)
+        return _get_event_args(charm, bound_event, self._juju_context)
 
-    def _emit_charm_event(self, event_name: str):
+    def _emit_charm_event(self, charm: _charm.CharmBase, event_name: str):
         """Emits a charm event based on a Juju event name.
 
         Args:
@@ -423,25 +451,52 @@ class _Manager:
             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)
+        event_to_emit = self._get_event_to_emit(charm, 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 = self._get_event_args(event_to_emit)
+        args, kwargs = self._get_event_args(charm, event_to_emit)
         logger.debug('Emitting Juju event %s.', event_name)
         event_to_emit.emit(*args, **kwargs)
 
-    def _commit(self):
+    def _commit(self, framework: _framework.Framework):
         """Commit the framework and gracefully teardown."""
-        self.framework.commit()
+        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_deferred(self):
+        """Emit deferred events and then commit the framework.
+
+        A framework and charm object are created for each notice in the storage
+        (an event and observer pair), the relevant deferred event is emitted,
+        and the framework is committed. Note that collect-status events are not
+        emitted.
+        """
+        # TODO: Remove the restricted context check below once we no longer need
+        #       to support Juju < 4 (collect-metrics and restricted context are
+        #       being removed in Juju 4.0).
+        #
+        # Skip re-emission of deferred events for collect-metrics events because
+        # they do not have the full access to all hook tools.
+        if self._dispatcher.is_restricted_context():
+            logger.debug('Skipping re-emission of deferred events in restricted context.')
+            return
+        # Re-emit previously deferred events to the observers that deferred them.
+        for event_path, _, _ in self._storage.notices():
+            event_handle = _framework.Handle.from_path(event_path)
+            logger.debug('Re-emitting deferred event: %s', event_handle)
+            charm = self._make_charm(event_handle.kind)
+            charm.framework._reemit_single_path(event_path)
+            self._commit(charm.framework)
+            self._close()
+            charm._destroy_charm()
+
     def _destroy(self):
         """Finalise the manager."""
         from . import tracing  # break circular import
@@ -453,11 +508,11 @@ class _Manager:
     def run(self):
         """Emit and then commit the framework."""
         try:
-            self._emit()
-            self._commit()
+            self._emit(self.charm, self._dispatcher.event_name)
+            self._commit(self.charm.framework)
             self._close()
         finally:
-            self.framework.close()
+            self.charm.framework.close()
 
 
 def main(charm_class: Type[_charm.CharmBase], use_juju_for_storage: Optional[bool] = None):
@@ -467,7 +522,10 @@ def main(charm_class: Type[_charm.CharmBase], use_juju_for_storage: Optional[boo
     """
     manager = None
     try:
-        manager = _Manager(charm_class, use_juju_for_storage=use_juju_for_storage)
+        juju_context = _JujuContext.from_dict(os.environ)
+        manager = _Manager(
+            charm_class, use_juju_for_storage=use_juju_for_storage, juju_context=juju_context
+        )
 
         manager.run()
     except _Abort as e:
diff --git a/ops/charm.py b/ops/charm.py
index ab5735a..7ee4edc 100644
--- a/ops/charm.py
+++ b/ops/charm.py
@@ -29,6 +29,7 @@ from typing import (
     Mapping,
     NoReturn,
     Optional,
+    Set,
     TextIO,
     Tuple,
     TypedDict,
@@ -1366,6 +1367,12 @@ class CharmBase(Object):
     def __init__(self, framework: Framework):
         super().__init__(framework, None)
 
+        # These events are present with the same names in all charms, unlike the
+        # ones that are defined below, where the presence and names depend on
+        # the relations, storages, actions, and containers defined in the charm
+        # metadata.
+        self._static_events: Set[str] = set(self.on.events())
+
         for relation_name in self.framework.meta.relations:
             relation_name = relation_name.replace('-', '_')
             self.on.define_event(f'{relation_name}_relation_created', RelationCreatedEvent)
@@ -1392,6 +1399,12 @@ class CharmBase(Object):
                 f'{container_name}_pebble_check_recovered', PebbleCheckRecoveredEvent
             )
 
+    def _destroy_charm(self):
+        for event in self.on.events():
+            if event in self._static_events:
+                continue
+            self.on._undefine_event(event)
+
     @property
     def app(self) -> model.Application:
         """Application that this unit is part of."""
diff --git a/ops/framework.py b/ops/framework.py
index c62af3f..4d84dcf 100644
--- a/ops/framework.py
+++ b/ops/framework.py
@@ -463,9 +463,9 @@ class ObjectEvents(Object):
         """
         prefix = 'unable to define an event with event_kind that '
         if not event_kind.isidentifier():
-            raise RuntimeError(f'{prefix}is not a valid python identifier: {event_kind}')
+            raise RuntimeError(f'{prefix}is not a valid Python identifier: {event_kind}')
         elif keyword.iskeyword(event_kind):
-            raise RuntimeError(f'{prefix}is a python keyword: {event_kind}')
+            raise RuntimeError(f'{prefix}is a Python keyword: {event_kind}')
         try:
             getattr(cls, event_kind)
             raise RuntimeError(
@@ -478,6 +478,23 @@ class ObjectEvents(Object):
         event_descriptor.__set_name__(cls, event_kind)
         setattr(cls, event_kind, event_descriptor)
 
+    @classmethod
+    def _undefine_event(cls, event_kind: str):
+        """Remove the definition of an event on this type at runtime.
+
+        This undoes the effect of :meth:`define_event`. This is not intended
+        for use by charm authors, but rather for use by the ops library itself.
+        """
+        event_descriptor = getattr(cls, event_kind)
+        if (framework := getattr(event_descriptor, 'framework', None)) is not None:
+            framework._unregister_type(
+                event_descriptor.event_type, event_descriptor.emitter, event_descriptor.event_kind
+            )
+        try:
+            delattr(cls, event_kind)
+        except AttributeError:
+            raise RuntimeError(f'no event with event_kind {event_kind} to undefine') from None
+
     def _event_kinds(self) -> List[str]:
         event_kinds: List[str] = []
         # We have to iterate over the class rather than instance to allow for properties which
@@ -725,6 +742,23 @@ class Framework(Object):
         self._type_registry[parent_path, kind_] = cls
         self._type_known.add(cls)
 
+    def _unregister_type(
+        self,
+        cls: Type[Serializable],
+        parent: Optional[Union['Handle', 'Object']],
+        kind: Optional[str] = None,
+    ):
+        """Unregister a type from a handle."""
+        parent_path: Optional[str] = None
+        if isinstance(parent, Object):
+            parent_path = parent.handle.path
+        elif isinstance(parent, Handle):
+            parent_path = parent.path
+
+        kind_: str = kind or cls.handle_kind
+        del self._type_registry[parent_path, kind_]
+        self._type_known.remove(cls)
+
     def _validate_snapshot_data(
         self, value: Union['StoredStateData', 'EventBase'], data: Dict[str, Any]
     ):
@@ -932,7 +966,7 @@ class Framework(Object):
         been first emitted won't be notified, as that would mean potentially observing
         events out of order.
         """
-        self._reemit()
+        self._reemit(emitting_deferred=True)
 
     @contextmanager
     def _event_context(self, event_name: str):
@@ -964,7 +998,10 @@ class Framework(Object):
 
         self._event_name = old_event_name
 
-    def _reemit(self, single_event_path: Optional[str] = None):
+    def _reemit_single_path(self, single_event_path: str):
+        self._reemit(single_event_path, emitting_deferred=True)
+
+    def _reemit(self, single_event_path: Optional[str] = None, emitting_deferred: bool = False):
         last_event_path = None
         deferred = True
         for event_path, observer_path, method_name in self._storage.notices(single_event_path):
@@ -979,6 +1016,7 @@ class Framework(Object):
             try:
                 event = self.load_snapshot(event_handle)
             except NoTypeError:
+                logger.debug('Skipping notice %s - cannot find event class.', event_path)
                 self._storage.drop_notice(event_path, observer_path, method_name)
                 continue
 
@@ -987,7 +1025,7 @@ class Framework(Object):
             observer = self._observer.get(observer_path)
 
             if observer:
-                if single_event_path is None:
+                if emitting_deferred:
                     logger.debug('Re-emitting deferred event %s.', event)
                 elif isinstance(event, LifecycleEvent):
                     # Ignore Lifecycle events: they are "private" and not interesting.
diff --git a/ops/model.py b/ops/model.py
index fa03223..6e440f2 100644
--- a/ops/model.py
+++ b/ops/model.py
@@ -3699,10 +3699,8 @@ class _ModelBackend:
     def application_version_set(self, version: str) -> None:
         self._run('application-version-set', '--', version)
 
-    @classmethod
-    def log_split(
-        cls, message: str, max_len: int = MAX_LOG_LINE_LEN
-    ) -> Generator[str, None, None]:
+    @staticmethod
+    def log_split(message: str, max_len: int = MAX_LOG_LINE_LEN) -> Generator[str, None, None]:
         """Helper to handle log messages that are potentially too long.
 
         This is a generator that splits a message string into multiple chunks if it is too long


Test patch:
diff --git a/test/charms/test_main/src/charm.py b/test/charms/test_main/src/charm.py
index 69dc366..ff9ad1d 100755
--- a/test/charms/test_main/src/charm.py
+++ b/test/charms/test_main/src/charm.py
@@ -26,6 +26,12 @@ sys.path.append('lib')
 
 logger = logging.getLogger()
 
+# Note for ops developers: juju-log doesn't go anywhere useful much of the time
+# during unit tests, and test_main failures that subprocess out are often
+# difficult to debug. Uncomment this line to get more informative errors when
+# running the tests.
+# logger.addHandler(logging.StreamHandler(sys.stderr))
+
 
 class CustomEvent(ops.EventBase):
     pass
@@ -123,6 +129,18 @@ class Charm(ops.CharmBase):
         if os.getenv('TRY_EXCEPTHOOK', False):
             raise RuntimeError('failing as requested')
 
+        for name, event in self.on.events().items():
+            if isinstance(event, ops.LifecycleEvent) or name == 'custom':
+                continue
+            self.framework.observe(event, self._on_any_event)
+        if hasattr(self, 'charm_attribute'):
+            raise RuntimeError('charm instance was reused')
+
+    def _on_any_event(self, event: ops.EventBase):
+        # Note that doing this is bad behaviour: we're doing it here to make
+        # sure that the value is *not* retained.
+        self.charm_attribute = str(event)
+
     def _on_install(self, event: ops.InstallEvent):
         self._stored.on_install.append(type(event).__name__)
         self._stored.observed_event_types.append(type(event).__name__)
diff --git a/test/test_main.py b/test/test_main.py
index 9254cc0..241a643 100644
--- a/test/test_main.py
+++ b/test/test_main.py
@@ -211,7 +211,7 @@ class TestDispatch:
                     ops.main(MyCharm)
 
         assert mock_charm_event.call_count == 1
-        return mock_charm_event.call_args[0][0]
+        return mock_charm_event.call_args[0][1]
 
     def test_with_dispatch(self):
         """With dispatch, dispatch is used."""
@@ -423,6 +423,7 @@ class _TestMain(abc.ABC):
         assert isinstance(state, ops.BoundStoredState)
         assert list(state.observed_event_types) == ['InstallEvent']
 
+        # The config-changed handler always defers.
         state = self._simulate_event(
             fake_script, EventSpec(ops.ConfigChangedEvent, 'config-changed')
         )
@@ -752,6 +753,13 @@ class _TestMain(abc.ABC):
 
         expected = [
             VERSION_LOGLINE,
+            [
+                'juju-log',
+                '--log-level',
+                'DEBUG',
+                '--',
+                'Skipping re-emission of deferred events in restricted context.',
+            ],
             ['juju-log', '--log-level', 'DEBUG', '--', 'Emitting Juju event collect_metrics.'],
             ['add-metric', '--labels', 'bar=4.2', 'foo=42'],
             ['is-leader', '--format=json'],
diff --git a/testing/src/scenario/_ops_main_mock.py b/testing/src/scenario/_ops_main_mock.py
index 57b3ece..6a09f0f 100644
--- a/testing/src/scenario/_ops_main_mock.py
+++ b/testing/src/scenario/_ops_main_mock.py
@@ -12,10 +12,12 @@ from typing import (
     Any,
     Dict,
     FrozenSet,
+    Generic,
     List,
     Sequence,
     Set,
     Tuple,
+    cast,
 )
 
 import ops
@@ -23,9 +25,10 @@ import ops.jujucontext
 import ops.storage
 
 from ops.framework import _event_regex
-from ops._main import _Dispatcher, _Manager
+from ops._main import _Manager
 from ops._main import logger as ops_logger
 
+from .context import Context
 from .errors import BadOwnerPath, NoObserverError
 from .logger import logger as scenario_logger
 from .mocking import _MockModelBackend
@@ -114,7 +117,7 @@ class UnitStateDB:
             db.save_snapshot(stored_state._handle_path, stored_state.content)
 
 
-class Ops(_Manager):
+class Ops(_Manager, Generic[CharmType]):
     """Class to manage stepping through ops setup, event emission and framework commit."""
 
     def __init__(
@@ -126,10 +129,14 @@ class Ops(_Manager):
         juju_context: ops.jujucontext._JujuContext,
     ):
         self.state = state
+        # This is the event passed to `run`, corresponding to the event that
+        # caused Juju to start the framework.
         self.event = event
         self.context = context
         self.charm_spec = charm_spec
         self.store = None
+        self._juju_context = juju_context
+        self.captured_events: List[ops.EventBase] = []
 
         try:
             import ops_tracing._mock  # break circular import
@@ -139,15 +146,7 @@ class Ops(_Manager):
         except ImportError:
             self._tracing_mock = 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)
+        super().__init__(charm_class=self.charm_spec.charm_type, juju_context=juju_context)
 
     def _load_charm_meta(self):
         metadata = (self._charm_root / 'metadata.yaml').read_text()
@@ -164,6 +163,21 @@ class Ops(_Manager):
 
         return ops.CharmMeta.from_yaml(metadata, actions_metadata, config_metadata)
 
+    def _make_framework(self, *args: Any, **kwargs: Any):
+        return CapturingFramework(*args, context=self.context, **kwargs)
+
+    def _make_model_backend(self):
+        # The event here is used in the context of the Juju event that caused
+        # the framework to start, so we pass in the original one, even though
+        # this backend might be used for a deferred event.
+        return _MockModelBackend(
+            state=self.state,
+            event=self.event,
+            context=self.context,
+            charm_spec=self.charm_spec,
+            juju_context=self._juju_context,
+        )
+
     def _setup_root_logging(self):
         # The warnings module captures this in _showwarning_orig, but we
         # shouldn't really be using a private method, so capture it ourselves as
@@ -178,7 +192,7 @@ class Ops(_Manager):
         # useful in a testing context, so we reset it here.
         sys.excepthook = sys.__excepthook__
 
-    def _make_storage(self, _: _Dispatcher):
+    def _make_storage(self):
         # TODO: add use_juju_for_storage support
         storage = ops.storage.SQLiteStorage(':memory:')
         logger.info('Copying input state to storage.')
@@ -186,7 +200,7 @@ class Ops(_Manager):
         self.store.apply_state(self.state)
         return storage
 
-    def _get_event_to_emit(self, event_name: str):
+    def _get_event_to_emit(self, charm: ops.CharmBase, event_name: str):
         if self.event._is_custom_event:
             assert self.event.custom_event is not None
             emitter_type = type(self.event.custom_event.emitter)
@@ -204,12 +218,12 @@ class Ops(_Manager):
                         continue
                     return getattr(sub_attr, self.event.custom_event.event_kind)
 
-        owner = self._get_owner(self.charm, self.event.owner_path) if self.event else self.charm.on
+        owner = self._get_owner(charm, self.event.owner_path) if self.event else 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)
+            ops_logger.debug('Event %s not defined for %s.', event_name, charm)
             raise NoObserverError(
                 f'Cannot fire {event_name!r} on {owner}: invalid event (not on charm.on).',
             )
@@ -231,14 +245,14 @@ class Ops(_Manager):
         return obj
 
     def _get_event_args(
-        self, bound_event: ops.framework.BoundEvent
+        self, charm: ops.CharmBase, bound_event: ops.framework.BoundEvent
     ) -> Tuple[List[Any], Dict[str, Any]]:
         # For custom events, if the caller provided us with explicit args, we
         # merge them with the Juju ones (to handle libraries subclassing the
         # Juju events). We also handle converting from Scenario to ops types,
         # since the test code typically won't be able to create the ops objects,
         # as a model is required for many.
-        args, kwargs = super()._get_event_args(bound_event)
+        args, kwargs = super()._get_event_args(charm, bound_event)
         if self.event.custom_event_args is not None:
             for arg in self.event.custom_event_args:
                 args.append(self._object_to_ops_object(arg))
@@ -276,3 +290,41 @@ class Ops(_Manager):
         super()._destroy()
         if self._tracing_mock:
             self._tracing_mock.__exit__(None, None, None)
+
+
+class CapturingFramework(ops.Framework):
+    def __init__(self, *args: Any, context: Context[CharmType], **kwargs: Any):
+        super().__init__(*args, **kwargs)
+        self._context = context
+
+    def _emit(self, event: ops.EventBase):
+        if self._context.capture_framework_events or not isinstance(
+            event,
+            (ops.PreCommitEvent, ops.CommitEvent, ops.CollectStatusEvent),
+        ):
+            # Dump/undump the event to ensure any custom attributes are (re)set by restore().
+            event.restore(event.snapshot())
+            self._context.emitted_events.append(event)
+
+        return super()._emit(event)
+
+    def _reemit_single_path(self, single_event_path: str):
+        if not self._context.capture_deferred_events:
+            return super()._reemit_single_path(single_event_path)
+
+        # Load all notices from storage as events.
+        for event_path, _, _ in self._storage.notices(single_event_path):
+            event_handle = ops.Handle.from_path(event_path)
+            try:
+                event = self.load_snapshot(event_handle)
+            except ops.NoTypeError:
+                continue
+            event = cast(ops.EventBase, event)
+            event.deferred = False
+            self._forget(event)  # prevent tracking conflicts
+
+            # Dump/undump the event to ensure any custom attributes are (re)set by restore().
+            event.restore(event.snapshot())
+            self._context.emitted_events.append(event)
+
+        return super()._reemit_single_path(single_event_path)
diff --git a/testing/src/scenario/_runtime.py b/testing/src/scenario/_runtime.py
index 43889e3..376d802 100644
--- a/testing/src/scenario/_runtime.py
+++ b/testing/src/scenario/_runtime.py
@@ -3,34 +3,23 @@
 
 """Test framework runtime."""
 
+import contextlib
 import copy
 import dataclasses
 import os
+import pathlib
 import tempfile
 import typing
-from contextlib import contextmanager
-from pathlib import Path
 from typing import (
     TYPE_CHECKING,
     Dict,
-    List,
     Optional,
     Type,
-    TypeVar,
     Union,
 )
 
 import yaml
-from ops import (
-    CollectStatusEvent,
-    pebble,
-    CommitEvent,
-    EventBase,
-    Framework,
-    Handle,
-    NoTypeError,
-    PreCommitEvent,
-)
+from ops import pebble
 from ops.jujucontext import _JujuContext
 from ops._private.harness import ActionFailed
 
@@ -48,7 +37,7 @@ if TYPE_CHECKING:  # pragma: no cover
 
 logger = scenario_logger.getChild('runtime')
 
-RUNTIME_MODULE = Path(__file__).parent
+RUNTIME_MODULE = pathlib.Path(__file__).parent
 
 
 class Runtime:
@@ -60,7 +49,7 @@ class Runtime:
     def __init__(
         self,
         charm_spec: '_CharmSpec[CharmType]',
-        charm_root: Optional[Union[str, Path]] = None,
+        charm_root: Optional[Union[str, pathlib.Path]] = None,
         juju_version: str = '3.0.0',
         app_name: Optional[str] = None,
         unit_id: Optional[int] = 0,
@@ -76,7 +65,7 @@ class Runtime:
         self._app_name = app_name
         self._unit_id = unit_id
 
-    def _get_event_env(self, state: 'State', event: '_Event', charm_root: Path):
+    def _get_event_env(self, state: 'State', event: '_Event', charm_root: pathlib.Path):
         """Build the simulated environment the operator framework expects."""
         env = {
             'JUJU_VERSION': self._juju_version,
@@ -202,7 +191,7 @@ class Runtime:
         WrappedCharm.__name__ = charm_type.__name__
         return typing.cast(Type['CharmType'], WrappedCharm)
 
-    @contextmanager
+    @contextlib.contextmanager
     def _virtual_charm_root(self):
         # If we are using runtime on a real charm, we can make some assumptions about the
         # directory structure we are going to find.
@@ -213,17 +202,17 @@ class Runtime:
 
         if charm_virtual_root := self._charm_root:
             charm_virtual_root_is_custom = True
-            virtual_charm_root = Path(charm_virtual_root)
+            virtual_charm_root = pathlib.Path(charm_virtual_root)
         else:
             charm_virtual_root = tempfile.TemporaryDirectory()
-            virtual_charm_root = Path(charm_virtual_root.name)
+            virtual_charm_root = pathlib.Path(charm_virtual_root.name)
             charm_virtual_root_is_custom = False
 
         metadata_yaml = virtual_charm_root / 'metadata.yaml'
         config_yaml = virtual_charm_root / 'config.yaml'
         actions_yaml = virtual_charm_root / 'actions.yaml'
 
-        metadata_files_present: Dict[Path, Optional[str]] = {
+        metadata_files_present: Dict[pathlib.Path, Optional[str]] = {
             file: file.read_text() if charm_virtual_root_is_custom and file.exists() else None
             for file in (metadata_yaml, config_yaml, actions_yaml)
         }
@@ -271,22 +260,18 @@ class Runtime:
             # charm_virtual_root is a tempdir
             typing.cast(tempfile.TemporaryDirectory, charm_virtual_root).cleanup()  # type: ignore
 
-    @contextmanager
-    def _exec_ctx(self, ctx: 'Context'):
+    @contextlib.contextmanager
+    def _exec_ctx(self, ctx: 'Context[CharmType]'):
         """python 3.8 compatibility shim"""
         with self._virtual_charm_root() as temporary_charm_root:
-            with capture_events(
-                include_deferred=ctx.capture_deferred_events,
-                include_framework=ctx.capture_framework_events,
-            ) as captured:
-                yield (temporary_charm_root, captured)
+            yield temporary_charm_root
 
-    @contextmanager
+    @contextlib.contextmanager
     def exec(
         self,
         state: 'State',
         event: '_Event',
-        context: 'Context',
+        context: 'Context[CharmType]',
     ):
         """Runs an event with this state as initial state on a charm.
 
@@ -307,7 +292,7 @@ class Runtime:
         output_state = copy.deepcopy(state)
 
         logger.info(' - generating virtual charm root')
-        with self._exec_ctx(context) as (temporary_charm_root, captured):
+        with self._exec_ctx(context) as temporary_charm_root:
             logger.info(' - preparing env')
             env = self._get_event_env(
                 state=state,
@@ -353,91 +338,5 @@ class Runtime:
                 os.environ.update(previous_env)
                 logger.info(' - exited ops.main')
 
-        context.emitted_events.extend(captured)
         logger.info('event dispatched. done.')
         context._set_output_state(ops.state)
-
-
-_T = TypeVar('_T', bound=EventBase)
-
-
-@contextmanager
-def capture_events(
-    *types: Type[EventBase],
-    include_framework: bool = False,
-    include_deferred: bool = True,
-):
-    """Capture all events of type `*types` (using instance checks).
-
-    Arguments exposed so that you can define your own fixtures if you want to.
-
-    Example::
-    >>> from ops import StartEvent
-    >>> from scenario import Event, State
-    >>> from charm import MyCustomEvent, MyCharm  # noqa
-    >>>
-    >>> def test_my_event():
-    >>>     with capture_events(StartEvent, MyCustomEvent) as captured:
-    >>>         trigger(State(), ("start", MyCharm, meta=MyCharm.META)
-    >>>
-    >>>     assert len(captured) == 2
-    >>>     e1, e2 = captured
-    >>>     assert isinstance(e2, MyCustomEvent)
-    >>>     assert e2.custom_attr == 'foo'
-    """
-    allowed_types = types or (EventBase,)
-
-    captured: List[EventBase] = []
-    _real_emit = Framework._emit
-    _real_reemit = Framework.reemit
-
-    def _wrapped_emit(self: Framework, evt: EventBase):
-        if not include_framework and isinstance(
-            evt,
-            (PreCommitEvent, CommitEvent, CollectStatusEvent),
-        ):
-            return _real_emit(self, evt)
-
-        if isinstance(evt, allowed_types):
-            # dump/undump the event to ensure any custom attributes are (re)set by restore()
-            evt.restore(evt.snapshot())
-            captured.append(evt)
-
-        return _real_emit(self, evt)
-
-    def _wrapped_reemit(self: Framework):
-        # Framework calls reemit() before emitting the main juju event. We intercept that call
-        # and capture all events in storage.
-
-        if not include_deferred:
-            return _real_reemit(self)
-
-        # load all notices from storage as events.
-        for event_path, _, _ in self._storage.notices():
-            event_handle = Handle.from_path(event_path)
-            try:
-                event = self.load_snapshot(event_handle)
-            except NoTypeError:
-                continue
-            event = typing.cast(EventBase, event)
-            event.deferred = False
-            self._forget(event)  # prevent tracking conflicts
-
-            if not include_framework and isinstance(
-                event,
-                (PreCommitEvent, CommitEvent),
-            ):
-                continue
-
-            if isinstance(event, allowed_types):
-                captured.append(event)
-
-        return _real_reemit(self)
-
-    Framework._emit = _wrapped_emit  # type: ignore
-    Framework.reemit = _wrapped_reemit
-
-    yield captured
-
-    Framework._emit = _real_emit
-    Framework.reemit = _real_reemit
diff --git a/testing/src/scenario/context.py b/testing/src/scenario/context.py
index 36dcde2..0f85edf 100644
--- a/testing/src/scenario/context.py
+++ b/testing/src/scenario/context.py
@@ -469,6 +469,11 @@ class Context(Generic[CharmType]):
             with ctx(ctx.on.start(), State()) as manager:
                 manager.charm._some_private_setup()
                 manager.run()
+
+    Note that the manager provides access to the charm that is used for the
+    primary event (the one passed as the first argument to ``Context``) but not
+    to the charm instances that are used to handle any events
+    in :attr:`State.deferred`.
     """
 
     juju_log: list[JujuLogLine]
diff --git a/testing/src/scenario/mocking.py b/testing/src/scenario/mocking.py
index b8639c7..d5db2ee 100644
--- a/testing/src/scenario/mocking.py
+++ b/testing/src/scenario/mocking.py
@@ -125,7 +125,7 @@ class _MockExecProcess:
             )
         return stdout, stderr
 
-    def send_signal(self, sig: Union[int, str]) -> NoReturn:  # noqa: U100
+    def send_signal(self, sig: Union[int, str]) -> NoReturn:
         """Send the given signal to the (mock) process."""
         raise NotImplementedError()
 
@@ -140,7 +140,7 @@ class _MockModelBackend(_ModelBackend):  # type: ignore
         state: 'State',
         event: '_Event',
         charm_spec: '_CharmSpec[CharmType]',
-        context: 'Context',
+        context: 'Context[CharmType]',
         juju_context: '_JujuContext',
     ):
         super().__init__(juju_context=juju_context)
@@ -196,7 +196,6 @@ class _MockModelBackend(_ModelBackend):  # type: ignore
                 container_root=container_root,
                 mounts=mounts,
                 state=self._state,
-                event=self._event,
                 charm_spec=self._charm_spec,
                 context=self._context,
                 container_name=container_name,
@@ -699,8 +698,8 @@ class _MockModelBackend(_ModelBackend):  # type: ignore
     # legacy ops API that we don't intend to mock:
     def pod_spec_set(
         self,
-        spec: Mapping[str, Any],  # noqa: U100
-        k8s_resources: Optional[Mapping[str, Any]] = None,  # noqa: U100
+        spec: Mapping[str, Any],
+        k8s_resources: Optional[Mapping[str, Any]] = None,
     ) -> NoReturn:
         raise NotImplementedError(
             'pod-spec-set is not implemented in Scenario (and probably never will be: '
@@ -709,8 +708,8 @@ class _MockModelBackend(_ModelBackend):  # type: ignore
 
     def add_metrics(
         self,
-        metrics: Mapping[str, Union[int, float]],  # noqa: U100
-        labels: Optional[Mapping[str, str]] = None,  # noqa: U100
+        metrics: Mapping[str, Union[int, float]],
+        labels: Optional[Mapping[str, str]] = None,
     ) -> NoReturn:
         raise NotImplementedError(
             'add-metrics is not implemented in Scenario (and probably never will be: '
@@ -750,14 +749,12 @@ class _MockPebbleClient(_TestingPebbleClient):
         mounts: Dict[str, Mount],
         *,
         state: 'State',
-        event: '_Event',
         charm_spec: '_CharmSpec[CharmType]',
-        context: 'Context',
+        context: 'Context[CharmType]',
         container_name: str,
     ):
         self._state = state
         self.socket_path = socket_path
-        self._event = event
         self._charm_spec = charm_spec
         self._context = context
         self._container_name = container_name
@@ -817,6 +814,7 @@ class _MockPebbleClient(_TestingPebbleClient):
                     spawn_time=now,
                     ready_time=now,
                 )
+                assert check.change_id is not None
                 self._changes[check.change_id] = change
 
     def get_plan(self) -> pebble.Plan:
diff --git a/testing/src/scenario/state.py b/testing/src/scenario/state.py
index d236705..6ad1734 100644
--- a/testing/src/scenario/state.py
+++ b/testing/src/scenario/state.py
@@ -2029,7 +2029,11 @@ class _Event:  # type: ignore
         return self._path.type is not _EventType.CUSTOM
 
     def deferred(self, handler: Callable[..., Any], event_id: int = 1) -> DeferredEvent:
-        """Construct a DeferredEvent from this Event."""
+        """Construct a DeferredEvent from this Event.
+
+        The charm class passed to the :class:`Context` must be observing the
+        event this DeferredEvent is based on.
+        """
         handler_repr = repr(handler)
         handler_re = re.compile(r'<function (.*) at .*>')
         match = handler_re.match(handler_repr)
diff --git a/testing/tests/test_e2e/test_deferred.py b/testing/tests/test_e2e/test_deferred.py
index 643d96e..b9c69b1 100644
--- a/testing/tests/test_e2e/test_deferred.py
+++ b/testing/tests/test_e2e/test_deferred.py
@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import logging
 import typing
 
 import ops
@@ -256,3 +257,27 @@ def test_defer_custom_event(mycharm):
         == {'arg0': 'foo', 'arg1': 28}
     )
     assert not state_2.deferred
+
+
+def test_juju_log_over_deferred(mycharm):
+    # A backend is created for each charm, so one for a deferred event and one
+    # for the primary event. We want to ensure that all the logs end up in the
+    # context.
+    logger = logging.getLogger(__name__)
+
+    class MyCharm(mycharm):
+        def _on_event(self, event: ops.EventBase):
+            super()._on_event(event)
+            logger.info('Handled event %s', event.__class__.__name__)
+
+    ctx = Context(MyCharm, meta=mycharm.META)
+    mycharm.defer_next = 1
+    ctx.run(
+        ctx.on.start(),
+        State(deferred=[ctx.on.update_status().deferred(MyCharm._on_event)]),
+    )
+    logs = [log.message for log in ctx.juju_log if log.level == 'INFO']
+    assert logs == [
+        'Handled event UpdateStatusEvent',
+        'Handled event StartEvent',
+    ]
diff --git a/testing/tests/test_e2e/test_event.py b/testing/tests/test_e2e/test_event.py
index 6e9f82d..abfbae7 100644
--- a/testing/tests/test_e2e/test_event.py
+++ b/testing/tests/test_e2e/test_event.py
@@ -64,7 +64,6 @@ def test_emitted_framework():
 
     ctx = Context(MyCharm, meta=MyCharm.META, capture_framework_events=True)
     ctx.run(ctx.on.update_status(), State())
-    assert len(ctx.emitted_events) == 4
     assert list(map(type, ctx.emitted_events)) == [
         ops.UpdateStatusEvent,
         ops.CollectStatusEvent,
@@ -95,9 +94,10 @@ def test_emitted_deferred():
         State(deferred=[ctx.on.update_status().deferred(MyCharm._on_update_status)]),
     )
 
-    assert len(ctx.emitted_events) == 5
     assert [e.handle.kind for e in ctx.emitted_events] == [
         'update_status',
+        'pre_commit',
+        'commit',
         'start',
         'collect_unit_status',
         'pre_commit',
diff --git a/testing/tests/test_e2e/test_juju_log.py b/testing/tests/test_e2e/test_juju_log.py
index bc0c15b..8fb3d5d 100644
--- a/testing/tests/test_e2e/test_juju_log.py
+++ b/testing/tests/test_e2e/test_juju_log.py
@@ -31,6 +31,6 @@ def mycharm():
 def test_juju_log(mycharm):
     ctx = Context(mycharm, meta=mycharm.META)
     ctx.run(ctx.on.start(), State())
-    assert ctx.juju_log[-2] == JujuLogLine(level='DEBUG', message='Emitting Juju event start.')
-    assert ctx.juju_log[-1] == JujuLogLine(level='WARNING', message='bar!')
+    assert JujuLogLine(level='DEBUG', message='Emitting Juju event start.') in ctx.juju_log
+    assert JujuLogLine(level='WARNING', message='bar!') in ctx.juju_log
     # prints are not juju-logged.
diff --git a/testing/tests/test_emitted_events_util.py b/testing/tests/test_emitted_events_util.py
index 7756727..3f17c30 100644
--- a/testing/tests/test_emitted_events_util.py
+++ b/testing/tests/test_emitted_events_util.py
@@ -1,83 +1,76 @@
-from ops.charm import CharmBase, CharmEvents, CollectStatusEvent, StartEvent
-from ops.framework import CommitEvent, EventBase, EventSource, PreCommitEvent
+from typing import Any
 
-from scenario import State
+import ops
+
+from scenario import Context, State
 from scenario.state import _Event
-from scenario._runtime import capture_events
-from .helpers import trigger
 
 
-class Foo(EventBase):
+class Foo(ops.EventBase):
     pass
 
 
-class MyCharmEvents(CharmEvents):
-    foo = EventSource(Foo)
+class MyCharmEvents(ops.CharmEvents):
+    foo = ops.EventSource(Foo)
 
 
-class MyCharm(CharmBase):
+class MyCharm(ops.CharmBase):
     META = {'name': 'mycharm'}
-    on = MyCharmEvents()
+    on = MyCharmEvents()  # type: ignore
 
-    def __init__(self, *args, **kwargs):
+    def __init__(self, *args: Any, **kwargs: Any):
         super().__init__(*args, **kwargs)
         self.framework.observe(self.on.start, self._on_start)
         self.framework.observe(self.on.foo, self._on_foo)
 
-    def _on_start(self, e):
+    def _on_start(self, _: ops.StartEvent):
         self.on.foo.emit()
 
-    def _on_foo(self, e):
+    def _on_foo(self, _: Foo):
         pass
 
 
 def test_capture_custom_evt_nonspecific_capture_include_fw_evts():
-    with capture_events(include_framework=True) as emitted:
-        trigger(State(), 'start', MyCharm, meta=MyCharm.META)
+    ctx = Context(MyCharm, meta=MyCharm.META, capture_framework_events=True)
+    ctx.run(ctx.on.start(), State())
 
+    emitted = ctx.emitted_events
     assert len(emitted) == 5
-    assert isinstance(emitted[0], StartEvent)
+    assert isinstance(emitted[0], ops.StartEvent)
     assert isinstance(emitted[1], Foo)
-    assert isinstance(emitted[2], CollectStatusEvent)
-    assert isinstance(emitted[3], PreCommitEvent)
-    assert isinstance(emitted[4], CommitEvent)
+    assert isinstance(emitted[2], ops.CollectStatusEvent)
+    assert isinstance(emitted[3], ops.PreCommitEvent)
+    assert isinstance(emitted[4], ops.CommitEvent)
 
 
 def test_capture_juju_evt():
-    with capture_events() as emitted:
-        trigger(State(), 'start', MyCharm, meta=MyCharm.META)
+    ctx = Context(MyCharm, meta=MyCharm.META)
+    ctx.run(ctx.on.start(), State())
 
+    emitted = ctx.emitted_events
     assert len(emitted) == 2
-    assert isinstance(emitted[0], StartEvent)
+    assert isinstance(emitted[0], ops.StartEvent)
     assert isinstance(emitted[1], Foo)
 
 
 def test_capture_deferred_evt():
-    # todo: this test should pass with ops < 2.1 as well
-    with capture_events() as emitted:
-        trigger(
-            State(deferred=[_Event('foo').deferred(handler=MyCharm._on_foo)]),
-            'start',
-            MyCharm,
-            meta=MyCharm.META,
-        )
+    ctx = Context(MyCharm, meta=MyCharm.META, capture_deferred_events=True)
+    deferred = [_Event('foo').deferred(handler=MyCharm._on_foo)]
+    ctx.run(ctx.on.start(), State(deferred=deferred))
 
+    emitted = ctx.emitted_events
     assert len(emitted) == 3
     assert isinstance(emitted[0], Foo)
-    assert isinstance(emitted[1], StartEvent)
+    assert isinstance(emitted[1], ops.StartEvent)
     assert isinstance(emitted[2], Foo)
 
 
 def test_capture_no_deferred_evt():
-    # todo: this test should pass with ops < 2.1 as well
-    with capture_events(include_deferred=False) as emitted:
-        trigger(
-            State(deferred=[_Event('foo').deferred(handler=MyCharm._on_foo)]),
-            'start',
-            MyCharm,
-            meta=MyCharm.META,
-        )
+    ctx = Context(MyCharm, meta=MyCharm.META)
+    deferred = [_Event('foo').deferred(handler=MyCharm._on_foo)]
+    ctx.run(ctx.on.start(), State(deferred=deferred))
 
+    emitted = ctx.emitted_events
     assert len(emitted) == 2
-    assert isinstance(emitted[0], StartEvent)
+    assert isinstance(emitted[0], ops.StartEvent)
     assert isinstance(emitted[1], Foo)