From 28948556d4a3bbc9a3129a6974895e451147043c Mon Sep 17 00:00:00 2001
From: Tony Meyer <tony.meyer@canonical.com>
Date: Thu, 8 Aug 2024 15:12:21 +1200
Subject: [PATCH] feat!: unify run() and run_action() and simplify context
 managers (#162)

The `run_action()` method (both standalone and in the context manager)
are removed. This means that all events, including action events, are
emitted with the same `run()` call, and both return the output state.

To get access to the results of actions, a new `action_output` attribute
is added to the `Context`. This is a simplified representation of the
Juju `operation` object (and the `task` objects in them), which are
displayed with `juju run`, but also available via `juju show-operation`.
The class has the same name as the Harness `ActionOutput` and will be
consolidated into a single class when Scenario is added to ops.

For example:

```python
# Scenario 6
out = ctx.run_action(Action("backup", params={"output": "data.tar.gz"}), State())
assert out.results == {"size": 1234}
assert out.logs = [..., ...]
assert out.state...

# Scenario 7
state = ctx.run(ctx.on.action("backup", params={"output": "data.tar.gz"}), State())
assert ctx.action_output.results == {"size": 1234}
assert ctx.action_output.logs = [..., ...]
assert state...
```

When the charm calls `event.fail()`, this raises an exception, in the
same way that Harness does. For example:

```python
# Scenario 6
out = ctx.run_action("bad-action", State())
assert out.failure == "couldn't do the thing"

# Scenario 7
with pytest.raises(ActionFailed) as exc_info:
    ctx.run(ctx.on.action("bad-action"), State())
assert exc_info.value.message == "couldn't do the thing"
```

The advantage of this is that tests that the action is successful do not
need to have `assert ctx.action_output.status != "failed"`, which is
easy to miss.

In addition, the `Context.manager` and `Context.action_manager` methods
are replaced by the ability to use the `Context` object itself as a
context manager.

For example:

```python
ctx = Context(MyCharm)
with ctx(ctx.on.start(), State()) as event:
    event.charm.prepare()
    state = event.run()
    assert event.charm...
```

The same code is also used (with `ctx.on.action()`) for action events.

Advantages:
* Slightly shorter code (no ".manager" or ".action_manager")
* Avoids naming complications with "manager", "action_manager" and the
various alternatives proposed in #115.

The `.output` property of the context manager is also removed. The
context manager will still handle running the event if it's not done
explicitly, but if that's the case then the output is not available. We
want to encourage people to explicitly run the event, not rely on the
automated behaviour - although I think it does make sense that it does
run, rather than raise or end in a weird situation where the event never
ran.

This replaces #115 and #118, being a combination of ideas/discussion
from both, plus incorporating the unification of run/run_action
discussed here, and the "action fail raises" discussed elsewhere.

Also, as drive-by changes while names are being made consistent:
* `_Port` becomes `Port`
* `_RelationBase` becomes (again) `RelationBase`

---------

Co-authored-by: PietroPasotti <pietro.pasotti@canonical.com>
---
 README.md                         |  56 +++++---
 docs/custom_conf.py               |   1 +
 docs/index.rst                    |   1 -
 scenario/__init__.py              |   8 +-
 scenario/context.py               | 213 +++++++++---------------------
 scenario/mocking.py               |   9 +-
 scenario/runtime.py               |   4 +-
 scenario/state.py                 |  49 ++++---
 tests/helpers.py                  |   2 +-
 tests/test_charm_spec_autoload.py |   2 +-
 tests/test_context.py             |  36 ++---
 tests/test_context_on.py          | 122 +++++++++--------
 tests/test_e2e/test_actions.py    |  99 ++++++++++----
 tests/test_e2e/test_cloud_spec.py |   6 +-
 tests/test_e2e/test_manager.py    |  24 ++--
 tests/test_e2e/test_network.py    |   6 +-
 tests/test_e2e/test_pebble.py     |  10 +-
 tests/test_e2e/test_ports.py      |   4 +-
 tests/test_e2e/test_relations.py  |   6 +-
 tests/test_e2e/test_resource.py   |   2 +-
 tests/test_e2e/test_secrets.py    |  54 ++++----
 tests/test_e2e/test_storage.py    |  14 +-
 tests/test_e2e/test_vroot.py      |   4 +-
 23 files changed, 355 insertions(+), 377 deletions(-)

diff --git a/README.md b/README.md
index 982a21ff..0e2332ff 100644
--- a/README.md
+++ b/README.md
@@ -496,7 +496,7 @@ If you want to, you can override any of these relation or extra-binding associat
 
 ```python
 state = scenario.State(networks={
-  scenario.Network("foo", [BindAddress([Address('192.0.2.1')])])
+  scenario.Network("foo", [scenario.BindAddress([scenario.Address('192.0.2.1')])])
 })
 ```
 
@@ -726,8 +726,8 @@ storage = scenario.Storage("foo")
 # Setup storage with some content:
 (storage.get_filesystem(ctx) / "myfile.txt").write_text("helloworld")
 
-with ctx.manager(ctx.on.update_status(), scenario.State(storages={storage})) as mgr:
-    foo = mgr.charm.model.storages["foo"][0]
+with ctx(ctx.on.update_status(), scenario.State(storages={storage})) as manager:
+    foo = manager.charm.model.storages["foo"][0]
     loc = foo.location
     path = loc / "myfile.txt"
     assert path.exists()
@@ -924,9 +924,9 @@ import pathlib
 
 ctx = scenario.Context(MyCharm, meta={'name': 'juliette', "resources": {"foo": {"type": "oci-image"}}})
 resource = scenario.Resource(name='foo', path='/path/to/resource.tar')
-with ctx.manager(ctx.on.start(), scenario.State(resources={resource})) as mgr:
+with ctx(ctx.on.start(), scenario.State(resources={resource})) as manager:
     # If the charm, at runtime, were to call self.model.resources.fetch("foo"), it would get '/path/to/resource.tar' back.
-    path = mgr.charm.model.resources.fetch('foo')
+    path = manager.charm.model.resources.fetch('foo')
     assert path == pathlib.Path('/path/to/resource.tar')
 ```
 
@@ -988,7 +988,6 @@ class MyVMCharm(ops.CharmBase):
 An action is a special sort of event, even though `ops` handles them almost identically.
 In most cases, you'll want to inspect the 'results' of an action, or whether it has failed or
 logged something while executing. Many actions don't have a direct effect on the output state.
-For this reason, the output state is less prominent in the return type of `Context.run_action`.
 
 How to test actions with scenario:
 
@@ -1000,18 +999,32 @@ def test_backup_action():
 
     # If you didn't declare do_backup in the charm's metadata, 
     # the `ConsistencyChecker` will slap you on the wrist and refuse to proceed.
-    out: scenario.ActionOutput = ctx.run_action(ctx.on.action("do_backup"), scenario.State())
+    state = ctx.run(ctx.on.action("do_backup"), scenario.State())
 
-    # You can assert action results, logs, failure using the ActionOutput interface:
-    assert out.logs == ['baz', 'qux']
-    
-    if out.success:
-      # If the action did not fail, we can read the results:
-      assert out.results == {'foo': 'bar'}
+    # You can assert on action results and logs using the context:
+    assert ctx.action_logs == ['baz', 'qux']
+    assert ctx.action_results == {'foo': 'bar'}
+```
+
+## Failing Actions
+
+If the charm code calls `event.fail()` to indicate that the action has failed,
+an `ActionFailed` exception will be raised. This avoids having to include
+success checks in every test where the action is successful.
+
+```python
+def test_backup_action_failed():
+    ctx = scenario.Context(MyCharm)
+
+    with pytest.raises(ActionFailed) as exc_info:
+        ctx.run(ctx.on.action("do_backup"), scenario.State())
+    assert exc_info.value.message == "sorry, couldn't do the backup"
+    # The state is also available if that's required:
+    assert exc_info.value.state.get_container(...)
 
-    else:
-      # If the action fails, we can read a failure message:
-      assert out.failure == 'boo-hoo'
+    # You can still assert action results and logs that occured as well as the failure:
+    assert ctx.action_logs == ['baz', 'qux']
+    assert ctx.action_results == {'foo': 'bar'}
 ```
 
 ## Parametrized Actions
@@ -1024,7 +1037,7 @@ def test_backup_action():
 
     # If the parameters (or their type) don't match what is declared in the metadata, 
     # the `ConsistencyChecker` will slap you on the other wrist.
-    out: scenario.ActionOutput = ctx.run_action(
+    state = ctx.run(
         ctx.on.action("do_backup", params={'a': 'b'}),
         scenario.State()
     )
@@ -1130,7 +1143,7 @@ Scenario is a black-box, state-transition testing framework. It makes it trivial
 B, but not to assert that, in the context of this charm execution, with this state, a certain charm-internal method was called and returned a
 given piece of data, or would return this and that _if_ it had been called.
 
-Scenario offers a cheekily-named context manager for this use case specifically:
+The Scenario `Context` object can be used as a context manager for this use case specifically:
 
 ```python notest
 from charms.bar.lib_name.v1.charm_lib import CharmLib
@@ -1152,8 +1165,7 @@ class MyCharm(ops.CharmBase):
 
 def test_live_charm_introspection(mycharm):
     ctx = scenario.Context(mycharm, meta=mycharm.META)
-    # If you want to do this with actions, you can use `Context.action_manager` instead.
-    with ctx.manager("start", scenario.State()) as manager:
+    with ctx(ctx.on.start(), scenario.State()) as manager:
         # This is your charm instance, after ops has set it up:
         charm: MyCharm = manager.charm
         
@@ -1174,8 +1186,8 @@ def test_live_charm_introspection(mycharm):
     assert state_out.unit_status == ...
 ```
 
-Note that you can't call `manager.run()` multiple times: the manager is a context that ensures that `ops.main` 'pauses' right
-before emitting the event to hand you some introspection hooks, but for the rest this is a regular scenario test: you
+Note that you can't call `manager.run()` multiple times: the object is a context that ensures that `ops.main` 'pauses' right
+before emitting the event to hand you some introspection hooks, but for the rest this is a regular Scenario test: you
 can't emit multiple events in a single charm execution.
 
 # The virtual charm root
diff --git a/docs/custom_conf.py b/docs/custom_conf.py
index 3e035474..10deb009 100644
--- a/docs/custom_conf.py
+++ b/docs/custom_conf.py
@@ -311,5 +311,6 @@ def _compute_navigation_tree(context):
     ('py:class', '_Event'),
     ('py:class', 'scenario.state._DCBase'),
     ('py:class', 'scenario.state._EntityStatus'),
+    ('py:class', 'scenario.state._Event'),
     ('py:class', 'scenario.state._max_posargs.<locals>._MaxPositionalArgs'),
 ]
diff --git a/docs/index.rst b/docs/index.rst
index 1a261b3d..4d1af4d9 100644
--- a/docs/index.rst
+++ b/docs/index.rst
@@ -17,7 +17,6 @@ scenario.Context
 
 .. automodule:: scenario.context
 
-
 scenario.consistency_checker
 ============================
 
diff --git a/scenario/__init__.py b/scenario/__init__.py
index bbd2c694..24c1cac0 100644
--- a/scenario/__init__.py
+++ b/scenario/__init__.py
@@ -1,8 +1,9 @@
 #!/usr/bin/env python3
 # Copyright 2023 Canonical Ltd.
 # See LICENSE file for licensing details.
-from scenario.context import ActionOutput, Context
+from scenario.context import Context, Manager
 from scenario.state import (
+    ActionFailed,
     ActiveStatus,
     Address,
     BindAddress,
@@ -21,6 +22,7 @@
     Network,
     Notice,
     PeerRelation,
+    Port,
     Relation,
     Resource,
     Secret,
@@ -37,7 +39,7 @@
 )
 
 __all__ = [
-    "ActionOutput",
+    "ActionFailed",
     "CheckInfo",
     "CloudCredential",
     "CloudSpec",
@@ -56,6 +58,7 @@
     "Address",
     "BindAddress",
     "Network",
+    "Port",
     "ICMPPort",
     "TCPPort",
     "UDPPort",
@@ -70,4 +73,5 @@
     "MaintenanceStatus",
     "ActiveStatus",
     "UnknownStatus",
+    "Manager",
 ]
diff --git a/scenario/context.py b/scenario/context.py
index d3efedc5..7998ceb4 100644
--- a/scenario/context.py
+++ b/scenario/context.py
@@ -1,7 +1,6 @@
 #!/usr/bin/env python3
 # Copyright 2023 Canonical Ltd.
 # See LICENSE file for licensing details.
-import dataclasses
 import tempfile
 from contextlib import contextmanager
 from pathlib import Path
@@ -12,6 +11,7 @@
 from scenario.logger import logger as scenario_logger
 from scenario.runtime import Runtime
 from scenario.state import (
+    ActionFailed,
     CheckInfo,
     Container,
     MetadataNotFoundError,
@@ -21,7 +21,6 @@
     _Action,
     _CharmSpec,
     _Event,
-    _max_posargs,
 )
 
 if TYPE_CHECKING:  # pragma: no cover
@@ -37,36 +36,12 @@
 DEFAULT_JUJU_VERSION = "3.4"
 
 
-@dataclasses.dataclass(frozen=True)
-class ActionOutput(_max_posargs(0)):
-    """Wraps the results of running an action event with ``run_action``."""
-
-    state: "State"
-    """The charm state after the action has been handled.
-
-    In most cases, actions are not expected to be affecting it."""
-    logs: List[str]
-    """Any logs associated with the action output, set by the charm with
-    :meth:`ops.ActionEvent.log`."""
-    results: Optional[Dict[str, Any]] = None
-    """Key-value mapping assigned by the charm as a result of the action.
-    Will be None if the charm never calls :meth:`ops.ActionEvent.set_results`."""
-    failure: Optional[str] = None
-    """None if the action was successful, otherwise the message the charm set with
-    :meth:`ops.ActionEvent.fail`."""
-
-    @property
-    def success(self) -> bool:
-        """True if this action was a success, False otherwise."""
-        return self.failure is None
-
-
 class InvalidEventError(RuntimeError):
-    """raised when something is wrong with the event passed to Context.run_*"""
+    """raised when something is wrong with the event passed to Context.run"""
 
 
 class InvalidActionError(InvalidEventError):
-    """raised when something is wrong with the action passed to Context.run_action"""
+    """raised when something is wrong with an action passed to Context.run"""
 
 
 class ContextSetupError(RuntimeError):
@@ -77,13 +52,22 @@ class AlreadyEmittedError(RuntimeError):
     """Raised when ``run()`` is called more than once."""
 
 
-class _Manager:
-    """Context manager to offer test code some runtime charm object introspection."""
+class Manager:
+    """Context manager to offer test code some runtime charm object introspection.
+
+    This class should not be instantiated directly: use a :class:`Context`
+    in a ``with`` statement instead, for example::
+
+        ctx = Context(MyCharm)
+        with ctx(ctx.on.start(), State()) as manager:
+            manager.charm.setup()
+            manager.run()
+    """
 
     def __init__(
         self,
         ctx: "Context",
-        arg: Union[str, _Action, _Event],
+        arg: _Event,
         state_in: "State",
     ):
         self._ctx = ctx
@@ -91,25 +75,21 @@ def __init__(
         self._state_in = state_in
 
         self._emitted: bool = False
-        self._run = None
 
         self.ops: Optional["Ops"] = None
-        self.output: Optional[Union["State", ActionOutput]] = None
+        self.output: Optional["State"] = None
 
     @property
     def charm(self) -> CharmBase:
         if not self.ops:
             raise RuntimeError(
-                "you should __enter__ this contextmanager before accessing this",
+                "you should __enter__ this context manager before accessing this",
             )
         return cast(CharmBase, self.ops.charm)
 
     @property
     def _runner(self):
-        raise NotImplementedError("override in subclass")
-
-    def _get_output(self):
-        raise NotImplementedError("override in subclass")
+        return self._ctx._run  # noqa
 
     def __enter__(self):
         self._wrapped_ctx = wrapped_ctx = self._runner(self._arg, self._state_in)
@@ -117,57 +97,29 @@ def __enter__(self):
         self.ops = ops
         return self
 
-    def run(self) -> Union[ActionOutput, "State"]:
+    def run(self) -> "State":
         """Emit the event and proceed with charm execution.
 
         This can only be done once.
         """
         if self._emitted:
-            raise AlreadyEmittedError("Can only context.manager.run() once.")
+            raise AlreadyEmittedError("Can only run once.")
         self._emitted = True
 
         # wrap up Runtime.exec() so that we can gather the output state
         self._wrapped_ctx.__exit__(None, None, None)
 
-        self.output = out = self._get_output()
-        return out
+        assert self._ctx._output_state is not None
+        return self._ctx._output_state
 
     def __exit__(self, exc_type, exc_val, exc_tb):  # noqa: U100
         if not self._emitted:
-            logger.debug("manager not invoked. Doing so implicitly...")
+            logger.debug(
+                "user didn't emit the event within the context manager scope. Doing so implicitly upon exit...",
+            )
             self.run()
 
 
-class _EventManager(_Manager):
-    if TYPE_CHECKING:  # pragma: no cover
-        output: State  # pyright: ignore[reportIncompatibleVariableOverride]
-
-        def run(self) -> "State":
-            return cast("State", super().run())
-
-    @property
-    def _runner(self):
-        return self._ctx._run_event  # noqa
-
-    def _get_output(self):
-        return self._ctx._output_state  # noqa
-
-
-class _ActionManager(_Manager):
-    if TYPE_CHECKING:  # pragma: no cover
-        output: ActionOutput  # pyright: ignore[reportIncompatibleVariableOverride]
-
-        def run(self) -> "ActionOutput":
-            return cast("ActionOutput", super().run())
-
-    @property
-    def _runner(self):
-        return self._ctx._run  # noqa
-
-    def _get_output(self):
-        return self._ctx._finalize_action(self._ctx.output_state)  # noqa
-
-
 class _CharmEvents:
     """Events generated by Juju pertaining to application lifecycle.
 
@@ -360,14 +312,12 @@ class Context:
     It contains: the charm source code being executed, the metadata files associated with it,
     a charm project repository root, and the Juju version to be simulated.
 
-    After you have instantiated ``Context``, typically you will call one of ``run()`` or
-    ``run_action()`` to execute the charm once, write any assertions you like on the output
-    state returned by the call, write any assertions you like on the ``Context`` attributes,
-    then discard the ``Context``.
+    After you have instantiated ``Context``, typically you will call ``run()``to execute the charm
+    once, write any assertions you like on the output state returned by the call, write any
+    assertions you like on the ``Context`` attributes, then discard the ``Context``.
 
     Each ``Context`` instance is in principle designed to be single-use:
     ``Context`` is not cleaned up automatically between charm runs.
-    You can call ``.clear()`` to do some clean up, but we don't guarantee all state will be gone.
 
     Any side effects generated by executing the charm, that are not rightful part of the
     ``State``, are in fact stored in the ``Context``:
@@ -378,6 +328,10 @@ class Context:
     - :attr:`workload_version_history`: record of the workload versions the charm has set
     - :attr:`removed_secret_revisions`: record of the secret revisions the charm has removed
     - :attr:`emitted_events`: record of the events (including custom) that the charm has processed
+    - :attr:`action_logs`: logs associated with the action output, set by the charm with
+        :meth:`ops.ActionEvent.log`
+    - :attr:`action_results`: key-value mapping assigned by the charm as a result of the action.
+        Will be None if the charm never calls :meth:`ops.ActionEvent.set_results`
 
     This allows you to write assertions not only on the output state, but also, to some
     extent, on the path the charm took to get there.
@@ -410,19 +364,16 @@ def test_foo():
         (local_path / 'foo' / 'bar.yaml').write_text('foo: bar')
         scenario.Context(... charm_root=virtual_root).run(...)
 
-    Args:
-        charm_type: the CharmBase subclass to call :meth:`ops.main` on.
-        meta: charm metadata to use. Needs to be a valid metadata.yaml format (as a dict).
-            If none is provided, we will search for a ``metadata.yaml`` file in the charm root.
-        actions: charm actions to use. Needs to be a valid actions.yaml format (as a dict).
-            If none is provided, we will search for a ``actions.yaml`` file in the charm root.
-        config: charm config to use. Needs to be a valid config.yaml format (as a dict).
-            If none is provided, we will search for a ``config.yaml`` file in the charm root.
-        juju_version: Juju agent version to simulate.
-        app_name: App name that this charm is deployed as. Defaults to the charm name as
-            defined in its metadata
-        unit_id: Unit ID that this charm is deployed as. Defaults to 0.
-        charm_root: virtual charm root the charm will be executed with.
+    If you need access to the charm object that will handle the event, use the
+    class in a ``with`` statement, like::
+
+        import scenario
+
+        def test_foo():
+            ctx = scenario.Context(MyCharm)
+            with ctx(ctx.on.start(), State()) as manager:
+                manager.charm._some_private_setup()
+                manager.run()
     """
 
     def __init__(
@@ -507,11 +458,10 @@ def __init__(
         # set by Runtime.exec() in self._run()
         self._output_state: Optional["State"] = None
 
-        # ephemeral side effects from running an action
-
-        self._action_logs: List[str] = []
-        self._action_results: Optional[Dict[str, str]] = None
-        self._action_failure: Optional[str] = None
+        # operations (and embedded tasks) from running actions
+        self.action_logs: List[str] = []
+        self.action_results: Optional[Dict[str, Any]] = None
+        self._action_failure_message: Optional[str] = None
 
         self.on = _CharmEvents()
 
@@ -550,14 +500,14 @@ def _record_status(self, state: "State", is_app: bool):
         else:
             self.unit_status_history.append(state.unit_status)
 
-    def manager(self, event: "_Event", state: "State"):
+    def __call__(self, event: "_Event", state: "State"):
         """Context manager to introspect live charm object before and after the event is emitted.
 
         Usage::
 
             ctx = Context(MyCharm)
-            with ctx.manager(ctx.on.start(), State()) as manager:
-                assert manager.charm._some_private_attribute == "foo"  # noqa
+            with ctx(ctx.on.start(), State()) as manager:
+                manager.charm._some_private_setup()
                 manager.run()  # this will fire the event
                 assert manager.charm._some_private_attribute == "bar"  # noqa
 
@@ -565,27 +515,7 @@ def manager(self, event: "_Event", state: "State"):
             event: the :class:`Event` that the charm will respond to.
             state: the :class:`State` instance to use when handling the Event.
         """
-        return _EventManager(self, event, state)
-
-    def action_manager(self, action: "_Action", state: "State"):
-        """Context manager to introspect live charm object before and after the event is emitted.
-
-        Usage:
-        >>> with Context().action_manager(Action("foo"), State()) as manager:
-        >>>     assert manager.charm._some_private_attribute == "foo"  # noqa
-        >>>     manager.run()  # this will fire the event
-        >>>     assert manager.charm._some_private_attribute == "bar"  # noqa
-
-        :arg action: the Action that the charm will execute.
-        :arg state: the State instance to use as data source for the hook tool calls that the
-            charm will invoke when handling the Action (event).
-        """
-        return _ActionManager(self, action, state)
-
-    @contextmanager
-    def _run_event(self, event: "_Event", state: "State"):
-        with self._run(event=event, state=state) as ops:
-            yield ops
+        return Manager(self, event, state)
 
     def run(self, event: "_Event", state: "State") -> "State":
         """Trigger a charm execution with an Event and a State.
@@ -597,40 +527,19 @@ def run(self, event: "_Event", state: "State") -> "State":
         :arg state: the State instance to use as data source for the hook tool calls that the
             charm will invoke when handling the Event.
         """
-        if isinstance(event, _Action) or event.action:
-            raise InvalidEventError("Use run_action() to run an action event.")
-        with self._run_event(event=event, state=state) as ops:
-            ops.emit()
-        return self.output_state
-
-    def run_action(self, event: "_Event", state: "State") -> ActionOutput:
-        """Trigger a charm execution with an action event and a State.
-
-        Calling this function will call ``ops.main`` and set up the context according to the
-        specified ``State``, then emit the event on the charm.
-
-        :arg event: the action event that the charm will execute.
-        :arg state: the State instance to use as data source for the hook tool calls that the
-            charm will invoke when handling the action event.
-        """
+        if event.action:
+            # Reset the logs, failure status, and results, in case the context
+            # is reused.
+            self.action_logs.clear()
+            if self.action_results is not None:
+                self.action_results.clear()
+            self._action_failure_message = None
         with self._run(event=event, state=state) as ops:
             ops.emit()
-        return self._finalize_action(self.output_state)
-
-    def _finalize_action(self, state_out: "State"):
-        ao = ActionOutput(
-            state=state_out,
-            logs=self._action_logs,
-            results=self._action_results,
-            failure=self._action_failure,
-        )
-
-        # reset all action-related state
-        self._action_logs = []
-        self._action_results = None
-        self._action_failure = None
-
-        return ao
+        if event.action:
+            if self._action_failure_message is not None:
+                raise ActionFailed(self._action_failure_message, self.output_state)
+        return self.output_state
 
     @contextmanager
     def _run(self, event: "_Event", state: "State"):
diff --git a/scenario/mocking.py b/scenario/mocking.py
index 0667cb7d..cf0480bf 100644
--- a/scenario/mocking.py
+++ b/scenario/mocking.py
@@ -551,21 +551,24 @@ def action_set(self, results: Dict[str, Any]):
         _format_action_result_dict(results)
         # but then we will store it in its unformatted,
         # original form for testing ease
-        self._context._action_results = results
+        if self._context.action_results:
+            self._context.action_results.update(results)
+        else:
+            self._context.action_results = results
 
     def action_fail(self, message: str = ""):
         if not self._event.action:
             raise ActionMissingFromContextError(
                 "not in the context of an action event: cannot action-fail",
             )
-        self._context._action_failure = message
+        self._context._action_failure_message = message
 
     def action_log(self, message: str):
         if not self._event.action:
             raise ActionMissingFromContextError(
                 "not in the context of an action event: cannot action-log",
             )
-        self._context._action_logs.append(message)
+        self._context.action_logs.append(message)
 
     def action_get(self):
         action = self._event.action
diff --git a/scenario/runtime.py b/scenario/runtime.py
index 2f739f8f..e853c682 100644
--- a/scenario/runtime.py
+++ b/scenario/runtime.py
@@ -20,7 +20,7 @@
 from scenario.capture_events import capture_events
 from scenario.logger import logger as scenario_logger
 from scenario.ops_main_mock import NoObserverError
-from scenario.state import DeferredEvent, PeerRelation, StoredState
+from scenario.state import ActionFailed, DeferredEvent, PeerRelation, StoredState
 
 if TYPE_CHECKING:  # pragma: no cover
     from ops.testing import CharmType
@@ -466,7 +466,7 @@ def exec(
                 # if the caller did not manually emit or commit: do that.
                 ops.finalize()
 
-            except NoObserverError:
+            except (NoObserverError, ActionFailed):
                 raise  # propagate along
             except Exception as e:
                 raise UncaughtCharmError(
diff --git a/scenario/state.py b/scenario/state.py
index 1bb190e2..789ddbfc 100644
--- a/scenario/state.py
+++ b/scenario/state.py
@@ -131,6 +131,14 @@ class MetadataNotFoundError(RuntimeError):
     """Raised when Scenario can't find a metadata.yaml file in the provided charm root."""
 
 
+class ActionFailed(Exception):
+    """Raised at the end of the hook if the charm has called `event.fail()`."""
+
+    def __init__(self, message: str, state: "State"):
+        self.message = message
+        self.state = state
+
+
 # This can be replaced with the KW_ONLY dataclasses functionality in Python 3.10+.
 def _max_posargs(n: int):
     class _MaxPositionalArgs:
@@ -429,7 +437,7 @@ def next_relation_id(*, update=True):
 
 
 @dataclasses.dataclass(frozen=True)
-class _RelationBase(_max_posargs(2)):
+class RelationBase(_max_posargs(2)):
     endpoint: str
     """Relation endpoint name. Must match some endpoint name defined in metadata.yaml."""
 
@@ -468,9 +476,9 @@ def _get_databag_for_remote(
         raise NotImplementedError()
 
     def __post_init__(self):
-        if type(self) is _RelationBase:
+        if type(self) is RelationBase:
             raise RuntimeError(
-                "_RelationBase cannot be instantiated directly; "
+                "RelationBase cannot be instantiated directly; "
                 "please use Relation, PeerRelation, or SubordinateRelation",
             )
 
@@ -502,7 +510,7 @@ def _validate_databag(self, databag: dict):
 
 
 @dataclasses.dataclass(frozen=True)
-class Relation(_RelationBase):
+class Relation(RelationBase):
     """An integration between the charm and another application."""
 
     remote_app_name: str = "remote"
@@ -546,7 +554,7 @@ def _databags(self):
 
 
 @dataclasses.dataclass(frozen=True)
-class SubordinateRelation(_RelationBase):
+class SubordinateRelation(RelationBase):
     remote_app_data: "RawDataBagContents" = dataclasses.field(default_factory=dict)
     remote_unit_data: "RawDataBagContents" = dataclasses.field(
         default_factory=lambda: DEFAULT_JUJU_DATABAG.copy(),
@@ -587,7 +595,7 @@ def remote_unit_name(self) -> str:
 
 
 @dataclasses.dataclass(frozen=True)
-class PeerRelation(_RelationBase):
+class PeerRelation(RelationBase):
     """A relation to share data between units of the charm."""
 
     peers_data: Dict["UnitID", "RawDataBagContents"] = dataclasses.field(
@@ -1079,8 +1087,12 @@ def __hash__(self) -> int:
 
 
 @dataclasses.dataclass(frozen=True)
-class _Port(_max_posargs(1)):
-    """Represents a port on the charm host."""
+class Port(_max_posargs(1)):
+    """Represents a port on the charm host.
+
+    Port objects should not be instantiated directly: use TCPPort, UDPPort, or
+    ICMPPort instead.
+    """
 
     port: Optional[int] = None
     """The port to open. Required for TCP and UDP; not allowed for ICMP."""
@@ -1088,20 +1100,20 @@ class _Port(_max_posargs(1)):
     """The protocol that data transferred over the port will use."""
 
     def __post_init__(self):
-        if type(self) is _Port:
+        if type(self) is Port:
             raise RuntimeError(
-                "_Port cannot be instantiated directly; "
+                "Port cannot be instantiated directly; "
                 "please use TCPPort, UDPPort, or ICMPPort",
             )
 
     def __eq__(self, other: object) -> bool:
-        if isinstance(other, (_Port, ops.Port)):
+        if isinstance(other, (Port, ops.Port)):
             return (self.protocol, self.port) == (other.protocol, other.port)
         return False
 
 
 @dataclasses.dataclass(frozen=True)
-class TCPPort(_Port):
+class TCPPort(Port):
     """Represents a TCP port on the charm host."""
 
     port: int
@@ -1118,7 +1130,7 @@ def __post_init__(self):
 
 
 @dataclasses.dataclass(frozen=True)
-class UDPPort(_Port):
+class UDPPort(Port):
     """Represents a UDP port on the charm host."""
 
     port: int
@@ -1135,7 +1147,7 @@ def __post_init__(self):
 
 
 @dataclasses.dataclass(frozen=True)
-class ICMPPort(_Port):
+class ICMPPort(Port):
     """Represents an ICMP port on the charm host."""
 
     protocol: _RawPortProtocolLiteral = "icmp"
@@ -1229,7 +1241,7 @@ class State(_max_posargs(0)):
     If a storage is not attached, omit it from this listing."""
 
     # we don't use sets to make json serialization easier
-    opened_ports: Iterable[_Port] = dataclasses.field(default_factory=frozenset)
+    opened_ports: Iterable[Port] = dataclasses.field(default_factory=frozenset)
     """Ports opened by juju on this charm."""
     leader: bool = False
     """Whether this charm has leadership."""
@@ -1275,7 +1287,7 @@ def __post_init__(self):
             else:
                 raise TypeError(f"Invalid status.{name}: {val!r}")
         normalised_ports = [
-            _Port(protocol=port.protocol, port=port.port)
+            Port(protocol=port.protocol, port=port.port)
             if isinstance(port, ops.Port)
             else port
             for port in self.opened_ports
@@ -1331,7 +1343,7 @@ def _update_status(
         # bypass frozen dataclass
         object.__setattr__(self, name, new_status)
 
-    def _update_opened_ports(self, new_ports: FrozenSet[_Port]):
+    def _update_opened_ports(self, new_ports: FrozenSet[Port]):
         """Update the current opened ports."""
         # bypass frozen dataclass
         object.__setattr__(self, "opened_ports", new_ports)
@@ -1832,10 +1844,11 @@ class _Action(_max_posargs(1)):
 
         def test_backup_action():
             ctx = scenario.Context(MyCharm)
-            out: scenario.ActionOutput = ctx.run_action(
+            state = ctx.run(
                 ctx.on.action('do_backup', params={'filename': 'foo'}),
                 scenario.State()
             )
+            assert ctx.action_results == ...
     """
 
     name: str
diff --git a/tests/helpers.py b/tests/helpers.py
index c8060d1c..82161c79 100644
--- a/tests/helpers.py
+++ b/tests/helpers.py
@@ -58,7 +58,7 @@ def trigger(
             event = getattr(ctx.on, event)(tuple(state.containers)[0])
         else:
             event = getattr(ctx.on, event)()
-    with ctx.manager(event, state=state) as mgr:
+    with ctx(event, state=state) as mgr:
         if pre_event:
             pre_event(mgr.charm)
         state_out = mgr.run()
diff --git a/tests/test_charm_spec_autoload.py b/tests/test_charm_spec_autoload.py
index fb738f87..57b93a31 100644
--- a/tests/test_charm_spec_autoload.py
+++ b/tests/test_charm_spec_autoload.py
@@ -162,6 +162,6 @@ def test_config_defaults(tmp_path, legacy):
     ) as charm:
         # this would fail if there were no 'cuddles' relation defined in meta
         ctx = Context(charm)
-        with ctx.manager(ctx.on.start(), State()) as mgr:
+        with ctx(ctx.on.start(), State()) as mgr:
             mgr.run()
             assert mgr.charm.config["foo"] is True
diff --git a/tests/test_context.py b/tests/test_context.py
index 2ca8b93a..361b4543 100644
--- a/tests/test_context.py
+++ b/tests/test_context.py
@@ -3,8 +3,8 @@
 import pytest
 from ops import CharmBase
 
-from scenario import ActionOutput, Context, State
-from scenario.state import _Action, _Event, next_action_id
+from scenario import Context, State
+from scenario.state import _Event, next_action_id
 
 
 class MyCharm(CharmBase):
@@ -36,8 +36,8 @@ def test_run_action():
 
     with patch.object(ctx, "_run") as p:
         ctx._output_state = "foo"  # would normally be set within the _run call scope
-        output = ctx.run_action(ctx.on.action("do-foo"), state)
-        assert output.state == "foo"
+        output = ctx.run(ctx.on.action("do-foo"), state)
+        assert output == "foo"
 
     assert p.called
     e = p.call_args.kwargs["event"]
@@ -53,26 +53,18 @@ def test_run_action():
 @pytest.mark.parametrize("unit_id", (1, 2, 42))
 def test_app_name(app_name, unit_id):
     ctx = Context(MyCharm, meta={"name": "foo"}, app_name=app_name, unit_id=unit_id)
-    with ctx.manager(ctx.on.start(), State()) as mgr:
+    with ctx(ctx.on.start(), State()) as mgr:
         assert mgr.charm.app.name == app_name
         assert mgr.charm.unit.name == f"{app_name}/{unit_id}"
 
 
-def test_action_output_no_positional_arguments():
-    with pytest.raises(TypeError):
-        ActionOutput(None, None)
-
-
-def test_action_output_no_results():
-    class MyCharm(CharmBase):
-        def __init__(self, framework):
-            super().__init__(framework)
-            framework.observe(self.on.act_action, self._on_act_action)
-
-        def _on_act_action(self, _):
-            pass
-
+def test_context_manager():
     ctx = Context(MyCharm, meta={"name": "foo"}, actions={"act": {}})
-    out = ctx.run_action(ctx.on.action("act"), State())
-    assert out.results is None
-    assert out.failure is None
+    state = State()
+    with ctx(ctx.on.start(), state) as mgr:
+        mgr.run()
+        assert mgr.charm.meta.name == "foo"
+
+    with ctx(ctx.on.action("act"), state) as mgr:
+        mgr.run()
+        assert mgr.charm.meta.name == "foo"
diff --git a/tests/test_context_on.py b/tests/test_context_on.py
index 8ddbf4d4..32759fd4 100644
--- a/tests/test_context_on.py
+++ b/tests/test_context_on.py
@@ -64,7 +64,7 @@ def test_simple_events(event_name, event_kind):
     ctx = scenario.Context(ContextCharm, meta=META, actions=ACTIONS)
     # These look like:
     #   ctx.run(ctx.on.install(), state)
-    with ctx.manager(getattr(ctx.on, event_name)(), scenario.State()) as mgr:
+    with ctx(getattr(ctx.on, event_name)(), scenario.State()) as mgr:
         mgr.run()
         assert len(mgr.charm.observed) == 2
         assert isinstance(mgr.charm.observed[1], ops.CollectStatusEvent)
@@ -93,13 +93,13 @@ def test_simple_secret_events(as_kwarg, event_name, event_kind, owner):
     else:
         args = (secret,)
         kwargs = {}
-    with ctx.manager(getattr(ctx.on, event_name)(*args, **kwargs), state_in) as mgr:
+    with ctx(getattr(ctx.on, event_name)(*args, **kwargs), state_in) as mgr:
         mgr.run()
         assert len(mgr.charm.observed) == 2
         assert isinstance(mgr.charm.observed[1], ops.CollectStatusEvent)
-        event = mgr.charm.observed[0]
-        assert isinstance(event, event_kind)
-        assert event.secret.id == secret.id
+        mgr = mgr.charm.observed[0]
+        assert isinstance(mgr, event_kind)
+        assert mgr.secret.id == secret.id
 
 
 @pytest.mark.parametrize(
@@ -121,14 +121,14 @@ def test_revision_secret_events(event_name, event_kind):
     #   ctx.run(ctx.on.secret_expired(secret=secret, revision=revision), state)
     # The secret and revision must always be passed because the same event name
     # is used for all secrets.
-    with ctx.manager(getattr(ctx.on, event_name)(secret, revision=42), state_in) as mgr:
+    with ctx(getattr(ctx.on, event_name)(secret, revision=42), state_in) as mgr:
         mgr.run()
         assert len(mgr.charm.observed) == 2
         assert isinstance(mgr.charm.observed[1], ops.CollectStatusEvent)
-        event = mgr.charm.observed[0]
-        assert isinstance(event, event_kind)
-        assert event.secret.id == secret.id
-        assert event.revision == 42
+        mgr = mgr.charm.observed[0]
+        assert isinstance(mgr, event_kind)
+        assert mgr.secret.id == secret.id
+        assert mgr.revision == 42
 
 
 @pytest.mark.parametrize("event_name", ["secret_expired", "secret_remove"])
@@ -157,42 +157,42 @@ def test_storage_events(event_name, event_kind):
     state_in = scenario.State(storages=[storage])
     # These look like:
     #   ctx.run(ctx.on.storage_attached(storage), state)
-    with ctx.manager(getattr(ctx.on, event_name)(storage), state_in) as mgr:
+    with ctx(getattr(ctx.on, event_name)(storage), state_in) as mgr:
         mgr.run()
         assert len(mgr.charm.observed) == 2
         assert isinstance(mgr.charm.observed[1], ops.CollectStatusEvent)
-        event = mgr.charm.observed[0]
-        assert isinstance(event, event_kind)
-        assert event.storage.name == storage.name
-        assert event.storage.index == storage.index
+        mgr = mgr.charm.observed[0]
+        assert isinstance(mgr, event_kind)
+        assert mgr.storage.name == storage.name
+        assert mgr.storage.index == storage.index
 
 
 def test_action_event_no_params():
     ctx = scenario.Context(ContextCharm, meta=META, actions=ACTIONS)
     # These look like:
-    #   ctx.run_action(ctx.on.action(action), state)
-    with ctx.action_manager(ctx.on.action("act"), scenario.State()) as mgr:
+    #   ctx.run(ctx.on.action(action_name), state)
+    with ctx(ctx.on.action("act"), scenario.State()) as mgr:
         mgr.run()
         assert len(mgr.charm.observed) == 2
         assert isinstance(mgr.charm.observed[1], ops.CollectStatusEvent)
-        event = mgr.charm.observed[0]
-        assert isinstance(event, ops.ActionEvent)
+        mgr = mgr.charm.observed[0]
+        assert isinstance(mgr, ops.ActionEvent)
 
 
 def test_action_event_with_params():
     ctx = scenario.Context(ContextCharm, meta=META, actions=ACTIONS)
     # These look like:
-    #   ctx.run_action(ctx.on.action(action=action), state)
+    #   ctx.run(ctx.on.action(action=action), state)
     # So that any parameters can be included and the ID can be customised.
     call_event = ctx.on.action("act", params={"param": "hello"})
-    with ctx.action_manager(call_event, scenario.State()) as mgr:
+    with ctx(call_event, scenario.State()) as mgr:
         mgr.run()
         assert len(mgr.charm.observed) == 2
         assert isinstance(mgr.charm.observed[1], ops.CollectStatusEvent)
-        event = mgr.charm.observed[0]
-        assert isinstance(event, ops.ActionEvent)
-        assert event.id == call_event.action.id
-        assert event.params["param"] == call_event.action.params["param"]
+        mgr = mgr.charm.observed[0]
+        assert isinstance(mgr, ops.ActionEvent)
+        assert mgr.id == call_event.action.id
+        assert mgr.params["param"] == call_event.action.params["param"]
 
 
 def test_pebble_ready_event():
@@ -201,13 +201,13 @@ def test_pebble_ready_event():
     state_in = scenario.State(containers=[container])
     # These look like:
     #   ctx.run(ctx.on.pebble_ready(container), state)
-    with ctx.manager(ctx.on.pebble_ready(container), state_in) as mgr:
+    with ctx(ctx.on.pebble_ready(container), state_in) as mgr:
         mgr.run()
         assert len(mgr.charm.observed) == 2
         assert isinstance(mgr.charm.observed[1], ops.CollectStatusEvent)
-        event = mgr.charm.observed[0]
-        assert isinstance(event, ops.PebbleReadyEvent)
-        assert event.workload.name == container.name
+        mgr = mgr.charm.observed[0]
+        assert isinstance(mgr, ops.PebbleReadyEvent)
+        assert mgr.workload.name == container.name
 
 
 @pytest.mark.parametrize("as_kwarg", [True, False])
@@ -230,15 +230,15 @@ def test_relation_app_events(as_kwarg, event_name, event_kind):
     else:
         args = (relation,)
         kwargs = {}
-    with ctx.manager(getattr(ctx.on, event_name)(*args, **kwargs), state_in) as mgr:
+    with ctx(getattr(ctx.on, event_name)(*args, **kwargs), state_in) as mgr:
         mgr.run()
         assert len(mgr.charm.observed) == 2
         assert isinstance(mgr.charm.observed[1], ops.CollectStatusEvent)
-        event = mgr.charm.observed[0]
-        assert isinstance(event, event_kind)
-        assert event.relation.id == relation.id
-        assert event.app.name == relation.remote_app_name
-        assert event.unit is None
+        mgr = mgr.charm.observed[0]
+        assert isinstance(mgr, event_kind)
+        assert mgr.relation.id == relation.id
+        assert mgr.app.name == relation.remote_app_name
+        assert mgr.unit is None
 
 
 def test_relation_complex_name():
@@ -247,14 +247,14 @@ def test_relation_complex_name():
     ctx = scenario.Context(ContextCharm, meta=meta, actions=ACTIONS)
     relation = scenario.Relation("foo-bar-baz")
     state_in = scenario.State(relations=[relation])
-    with ctx.manager(ctx.on.relation_created(relation), state_in) as mgr:
+    with ctx(ctx.on.relation_created(relation), state_in) as mgr:
         mgr.run()
         assert len(mgr.charm.observed) == 2
-        event = mgr.charm.observed[0]
-        assert isinstance(event, ops.RelationCreatedEvent)
-        assert event.relation.id == relation.id
-        assert event.app.name == relation.remote_app_name
-        assert event.unit is None
+        mgr = mgr.charm.observed[0]
+        assert isinstance(mgr, ops.RelationCreatedEvent)
+        assert mgr.relation.id == relation.id
+        assert mgr.app.name == relation.remote_app_name
+        assert mgr.unit is None
 
 
 @pytest.mark.parametrize("event_name", ["relation_created", "relation_broken"])
@@ -280,15 +280,15 @@ def test_relation_unit_events_default_unit(event_name, event_kind):
     # These look like:
     #   ctx.run(ctx.on.baz_relation_changed, state)
     # The unit is chosen automatically.
-    with ctx.manager(getattr(ctx.on, event_name)(relation), state_in) as mgr:
+    with ctx(getattr(ctx.on, event_name)(relation), state_in) as mgr:
         mgr.run()
         assert len(mgr.charm.observed) == 2
         assert isinstance(mgr.charm.observed[1], ops.CollectStatusEvent)
-        event = mgr.charm.observed[0]
-        assert isinstance(event, event_kind)
-        assert event.relation.id == relation.id
-        assert event.app.name == relation.remote_app_name
-        assert event.unit.name == "remote/1"
+        mgr = mgr.charm.observed[0]
+        assert isinstance(mgr, event_kind)
+        assert mgr.relation.id == relation.id
+        assert mgr.app.name == relation.remote_app_name
+        assert mgr.unit.name == "remote/1"
 
 
 @pytest.mark.parametrize(
@@ -306,17 +306,15 @@ def test_relation_unit_events(event_name, event_kind):
     state_in = scenario.State(relations=[relation])
     # These look like:
     #   ctx.run(ctx.on.baz_relation_changed(unit=unit_ordinal), state)
-    with ctx.manager(
-        getattr(ctx.on, event_name)(relation, remote_unit=2), state_in
-    ) as mgr:
+    with ctx(getattr(ctx.on, event_name)(relation, remote_unit=2), state_in) as mgr:
         mgr.run()
         assert len(mgr.charm.observed) == 2
         assert isinstance(mgr.charm.observed[1], ops.CollectStatusEvent)
-        event = mgr.charm.observed[0]
-        assert isinstance(event, event_kind)
-        assert event.relation.id == relation.id
-        assert event.app.name == relation.remote_app_name
-        assert event.unit.name == "remote/2"
+        mgr = mgr.charm.observed[0]
+        assert isinstance(mgr, event_kind)
+        assert mgr.relation.id == relation.id
+        assert mgr.app.name == relation.remote_app_name
+        assert mgr.unit.name == "remote/2"
 
 
 def test_relation_departed_event():
@@ -325,15 +323,15 @@ def test_relation_departed_event():
     state_in = scenario.State(relations=[relation])
     # These look like:
     #   ctx.run(ctx.on.baz_relation_departed(unit=unit_ordinal, departing_unit=unit_ordinal), state)
-    with ctx.manager(
+    with ctx(
         ctx.on.relation_departed(relation, remote_unit=2, departing_unit=1), state_in
     ) as mgr:
         mgr.run()
         assert len(mgr.charm.observed) == 2
         assert isinstance(mgr.charm.observed[1], ops.CollectStatusEvent)
-        event = mgr.charm.observed[0]
-        assert isinstance(event, ops.RelationDepartedEvent)
-        assert event.relation.id == relation.id
-        assert event.app.name == relation.remote_app_name
-        assert event.unit.name == "remote/2"
-        assert event.departing_unit.name == "remote/1"
+        mgr = mgr.charm.observed[0]
+        assert isinstance(mgr, ops.RelationDepartedEvent)
+        assert mgr.relation.id == relation.id
+        assert mgr.app.name == relation.remote_app_name
+        assert mgr.unit.name == "remote/2"
+        assert mgr.departing_unit.name == "remote/1"
diff --git a/tests/test_e2e/test_actions.py b/tests/test_e2e/test_actions.py
index b0668355..7b6d1727 100644
--- a/tests/test_e2e/test_actions.py
+++ b/tests/test_e2e/test_actions.py
@@ -3,8 +3,7 @@
 from ops.charm import ActionEvent, CharmBase
 from ops.framework import Framework
 
-from scenario import Context
-from scenario.context import InvalidEventError
+from scenario import ActionFailed, Context
 from scenario.state import State, _Action, next_action_id
 
 
@@ -34,14 +33,29 @@ def test_action_event(mycharm, baz_value):
             "foo": {"params": {"bar": {"type": "number"}, "baz": {"type": "boolean"}}}
         },
     )
-    ctx.run_action(ctx.on.action("foo", params={"baz": baz_value, "bar": 10}), State())
+    state = ctx.run(ctx.on.action("foo", params={"baz": baz_value, "bar": 10}), State())
 
+    assert isinstance(state, State)
     evt = ctx.emitted_events[0]
-
     assert evt.params["bar"] == 10
     assert evt.params["baz"] is baz_value
 
 
+def test_action_no_results():
+    class MyCharm(CharmBase):
+        def __init__(self, framework):
+            super().__init__(framework)
+            framework.observe(self.on.act_action, self._on_act_action)
+
+        def _on_act_action(self, _):
+            pass
+
+    ctx = Context(MyCharm, meta={"name": "foo"}, actions={"act": {}})
+    ctx.run(ctx.on.action("act"), State())
+    assert ctx.action_results is None
+    assert ctx.action_logs == []
+
+
 @pytest.mark.parametrize("res_value", ("one", 1, [2], ["bar"], (1,), {1, 2}))
 def test_action_event_results_invalid(mycharm, res_value):
     def handle_evt(charm: CharmBase, evt: ActionEvent):
@@ -51,19 +65,12 @@ def handle_evt(charm: CharmBase, evt: ActionEvent):
     mycharm._evt_handler = handle_evt
 
     ctx = Context(mycharm, meta={"name": "foo"}, actions={"foo": {}})
-    ctx.run_action(ctx.on.action("foo"), State())
-
-
-def test_cannot_run_action(mycharm):
-    ctx = Context(mycharm, meta={"name": "foo"}, actions={"foo": {}})
-
-    with pytest.raises(InvalidEventError):
-        ctx.run(ctx.on.action("foo"), state=State())
+    ctx.run(ctx.on.action("foo"), State())
 
 
 @pytest.mark.parametrize("res_value", ({"a": {"b": {"c"}}}, {"d": "e"}))
 def test_action_event_results_valid(mycharm, res_value):
-    def handle_evt(charm: CharmBase, evt):
+    def handle_evt(_: CharmBase, evt):
         if not isinstance(evt, ActionEvent):
             return
         evt.set_results(res_value)
@@ -74,15 +81,14 @@ def handle_evt(charm: CharmBase, evt):
 
     ctx = Context(mycharm, meta={"name": "foo"}, actions={"foo": {}})
 
-    out = ctx.run_action(ctx.on.action("foo"), State())
+    ctx.run(ctx.on.action("foo"), State())
 
-    assert out.results == res_value
-    assert out.success is True
+    assert ctx.action_results == res_value
 
 
 @pytest.mark.parametrize("res_value", ({"a": {"b": {"c"}}}, {"d": "e"}))
 def test_action_event_outputs(mycharm, res_value):
-    def handle_evt(charm: CharmBase, evt: ActionEvent):
+    def handle_evt(_: CharmBase, evt: ActionEvent):
         if not isinstance(evt, ActionEvent):
             return
 
@@ -94,11 +100,31 @@ def handle_evt(charm: CharmBase, evt: ActionEvent):
     mycharm._evt_handler = handle_evt
 
     ctx = Context(mycharm, meta={"name": "foo"}, actions={"foo": {}})
-    out = ctx.run_action(ctx.on.action("foo"), State())
+    with pytest.raises(ActionFailed) as exc_info:
+        ctx.run(ctx.on.action("foo"), State())
+    assert exc_info.value.message == "failed becozz"
+    assert ctx.action_results == {"my-res": res_value}
+    assert ctx.action_logs == ["log1", "log2"]
+
+
+def test_action_continues_after_fail():
+    class MyCharm(CharmBase):
+        def __init__(self, framework):
+            super().__init__(framework)
+            framework.observe(self.on.foo_action, self._on_foo_action)
+
+        def _on_foo_action(self, event):
+            event.log("starting")
+            event.set_results({"initial": "result"})
+            event.fail("oh no!")
+            event.set_results({"final": "result"})
 
-    assert out.failure == "failed becozz"
-    assert out.logs == ["log1", "log2"]
-    assert out.success is False
+    ctx = Context(MyCharm, meta={"name": "foo"}, actions={"foo": {}})
+    with pytest.raises(ActionFailed) as exc_info:
+        ctx.run(ctx.on.action("foo"), State())
+    assert exc_info.value.message == "oh no!"
+    assert ctx.action_logs == ["starting"]
+    assert ctx.action_results == {"initial": "result", "final": "result"}
 
 
 def _ops_less_than(wanted_major, wanted_minor):
@@ -114,7 +140,7 @@ def _ops_less_than(wanted_major, wanted_minor):
     _ops_less_than(2, 11), reason="ops 2.10 and earlier don't have ActionEvent.id"
 )
 def test_action_event_has_id(mycharm):
-    def handle_evt(charm: CharmBase, evt: ActionEvent):
+    def handle_evt(_: CharmBase, evt: ActionEvent):
         if not isinstance(evt, ActionEvent):
             return
         assert isinstance(evt.id, str) and evt.id != ""
@@ -122,7 +148,7 @@ def handle_evt(charm: CharmBase, evt: ActionEvent):
     mycharm._evt_handler = handle_evt
 
     ctx = Context(mycharm, meta={"name": "foo"}, actions={"foo": {}})
-    ctx.run_action(ctx.on.action("foo"), State())
+    ctx.run(ctx.on.action("foo"), State())
 
 
 @pytest.mark.skipif(
@@ -139,7 +165,32 @@ def handle_evt(charm: CharmBase, evt: ActionEvent):
     mycharm._evt_handler = handle_evt
 
     ctx = Context(mycharm, meta={"name": "foo"}, actions={"foo": {}})
-    ctx.run_action(ctx.on.action("foo", id=uuid), State())
+    ctx.run(ctx.on.action("foo", id=uuid), State())
+
+
+def test_two_actions_same_context():
+    class MyCharm(CharmBase):
+        def __init__(self, framework):
+            super().__init__(framework)
+            framework.observe(self.on.foo_action, self._on_foo_action)
+            framework.observe(self.on.bar_action, self._on_bar_action)
+
+        def _on_foo_action(self, event):
+            event.log("foo")
+            event.set_results({"foo": "result"})
+
+        def _on_bar_action(self, event):
+            event.log("bar")
+            event.set_results({"bar": "result"})
+
+    ctx = Context(MyCharm, meta={"name": "foo"}, actions={"foo": {}, "bar": {}})
+    ctx.run(ctx.on.action("foo"), State())
+    assert ctx.action_results == {"foo": "result"}
+    assert ctx.action_logs == ["foo"]
+    # Not recommended, but run another action in the same context.
+    ctx.run(ctx.on.action("bar"), State())
+    assert ctx.action_results == {"bar": "result"}
+    assert ctx.action_logs == ["bar"]
 
 
 def test_positional_arguments():
diff --git a/tests/test_e2e/test_cloud_spec.py b/tests/test_e2e/test_cloud_spec.py
index 1834b3da..c3a09248 100644
--- a/tests/test_e2e/test_cloud_spec.py
+++ b/tests/test_e2e/test_cloud_spec.py
@@ -47,14 +47,14 @@ def test_get_cloud_spec():
             name="lxd-model", type="lxd", cloud_spec=scenario_cloud_spec
         ),
     )
-    with ctx.manager(ctx.on.start(), state=state) as mgr:
+    with ctx(ctx.on.start(), state=state) as mgr:
         assert mgr.charm.model.get_cloud_spec() == expected_cloud_spec
 
 
 def test_get_cloud_spec_error():
     ctx = scenario.Context(MyCharm, meta={"name": "foo"})
     state = scenario.State(model=scenario.Model(name="lxd-model", type="lxd"))
-    with ctx.manager(ctx.on.start(), state) as mgr:
+    with ctx(ctx.on.start(), state) as mgr:
         with pytest.raises(ops.ModelError):
             mgr.charm.model.get_cloud_spec()
 
@@ -65,6 +65,6 @@ def test_get_cloud_spec_untrusted():
     state = scenario.State(
         model=scenario.Model(name="lxd-model", type="lxd", cloud_spec=cloud_spec),
     )
-    with ctx.manager(ctx.on.start(), state) as mgr:
+    with ctx(ctx.on.start(), state) as mgr:
         with pytest.raises(ops.ModelError):
             mgr.charm.model.get_cloud_spec()
diff --git a/tests/test_e2e/test_manager.py b/tests/test_e2e/test_manager.py
index 3f99ffd0..1856c7ae 100644
--- a/tests/test_e2e/test_manager.py
+++ b/tests/test_e2e/test_manager.py
@@ -5,7 +5,7 @@
 from ops.charm import CharmBase, CollectStatusEvent
 
 from scenario import Context, State
-from scenario.context import ActionOutput, AlreadyEmittedError, _EventManager
+from scenario.context import AlreadyEmittedError, Manager
 
 
 @pytest.fixture(scope="function")
@@ -23,7 +23,6 @@ def _on_event(self, e):
             if isinstance(e, CollectStatusEvent):
                 return
 
-            print("event!")
             self.unit.status = ActiveStatus(e.handle.kind)
 
     return MyCharm
@@ -31,41 +30,34 @@ def _on_event(self, e):
 
 def test_manager(mycharm):
     ctx = Context(mycharm, meta=mycharm.META)
-    with _EventManager(ctx, ctx.on.start(), State()) as manager:
+    with Manager(ctx, ctx.on.start(), State()) as manager:
         assert isinstance(manager.charm, mycharm)
-        assert not manager.output
         state_out = manager.run()
-        assert manager.output is state_out
 
     assert isinstance(state_out, State)
-    assert manager.output  # still there!
 
 
 def test_manager_implicit(mycharm):
     ctx = Context(mycharm, meta=mycharm.META)
-    with _EventManager(ctx, ctx.on.start(), State()) as manager:
+    with Manager(ctx, ctx.on.start(), State()) as manager:
         assert isinstance(manager.charm, mycharm)
         # do not call .run()
 
     # run is called automatically
     assert manager._emitted
-    assert manager.output
-    assert manager.output.unit_status == ActiveStatus("start")
 
 
 def test_manager_reemit_fails(mycharm):
     ctx = Context(mycharm, meta=mycharm.META)
-    with _EventManager(ctx, ctx.on.start(), State()) as manager:
+    with Manager(ctx, ctx.on.start(), State()) as manager:
         manager.run()
         with pytest.raises(AlreadyEmittedError):
             manager.run()
 
-    assert manager.output
-
 
 def test_context_manager(mycharm):
     ctx = Context(mycharm, meta=mycharm.META)
-    with ctx.manager(ctx.on.start(), State()) as manager:
+    with ctx(ctx.on.start(), State()) as manager:
         state_out = manager.run()
         assert isinstance(state_out, State)
     assert ctx.emitted_events[0].handle.kind == "start"
@@ -73,7 +65,7 @@ def test_context_manager(mycharm):
 
 def test_context_action_manager(mycharm):
     ctx = Context(mycharm, meta=mycharm.META, actions=mycharm.ACTIONS)
-    with ctx.action_manager(ctx.on.action("do-x"), State()) as manager:
-        ao = manager.run()
-        assert isinstance(ao, ActionOutput)
+    with ctx(ctx.on.action("do-x"), State()) as manager:
+        state_out = manager.run()
+        assert isinstance(state_out, State)
     assert ctx.emitted_events[0].handle.kind == "do_x_action"
diff --git a/tests/test_e2e/test_network.py b/tests/test_e2e/test_network.py
index 5cceac92..f152e6b6 100644
--- a/tests/test_e2e/test_network.py
+++ b/tests/test_e2e/test_network.py
@@ -47,7 +47,7 @@ def test_ip_get(mycharm):
         },
     )
 
-    with ctx.manager(
+    with ctx(
         ctx.on.update_status(),
         State(
             relations=[
@@ -84,7 +84,7 @@ def test_no_sub_binding(mycharm):
         },
     )
 
-    with ctx.manager(
+    with ctx(
         ctx.on.update_status(),
         State(
             relations=[
@@ -109,7 +109,7 @@ def test_no_relation_error(mycharm):
         },
     )
 
-    with ctx.manager(
+    with ctx(
         ctx.on.update_status(),
         State(
             relations=[
diff --git a/tests/test_e2e/test_pebble.py b/tests/test_e2e/test_pebble.py
index da40cf5d..dec93c4d 100644
--- a/tests/test_e2e/test_pebble.py
+++ b/tests/test_e2e/test_pebble.py
@@ -128,7 +128,7 @@ def callback(self: CharmBase):
         charm_type=charm_cls,
         meta={"name": "foo", "containers": {"foo": {}}},
     )
-    with ctx.manager(ctx.on.start(), state=state) as mgr:
+    with ctx(ctx.on.start(), state=state) as mgr:
         out = mgr.run()
         callback(mgr.charm)
 
@@ -318,7 +318,7 @@ def test_exec_wait_error(charm_cls):
     )
 
     ctx = Context(charm_cls, meta={"name": "foo", "containers": {"foo": {}}})
-    with ctx.manager(ctx.on.start(), state) as mgr:
+    with ctx(ctx.on.start(), state) as mgr:
         container = mgr.charm.unit.get_container("foo")
         proc = container.exec(["foo"])
         with pytest.raises(ExecError):
@@ -340,7 +340,7 @@ def test_exec_wait_output(charm_cls):
     )
 
     ctx = Context(charm_cls, meta={"name": "foo", "containers": {"foo": {}}})
-    with ctx.manager(ctx.on.start(), state) as mgr:
+    with ctx(ctx.on.start(), state) as mgr:
         container = mgr.charm.unit.get_container("foo")
         proc = container.exec(["foo"])
         out, err = proc.wait_output()
@@ -360,7 +360,7 @@ def test_exec_wait_output_error(charm_cls):
     )
 
     ctx = Context(charm_cls, meta={"name": "foo", "containers": {"foo": {}}})
-    with ctx.manager(ctx.on.start(), state) as mgr:
+    with ctx(ctx.on.start(), state) as mgr:
         container = mgr.charm.unit.get_container("foo")
         proc = container.exec(["foo"])
         with pytest.raises(ExecError):
@@ -381,7 +381,7 @@ def test_pebble_custom_notice(charm_cls):
 
     state = State(containers=[container])
     ctx = Context(charm_cls, meta={"name": "foo", "containers": {"foo": {}}})
-    with ctx.manager(
+    with ctx(
         ctx.on.pebble_custom_notice(container=container, notice=notices[-1]), state
     ) as mgr:
         container = mgr.charm.unit.get_container("foo")
diff --git a/tests/test_e2e/test_ports.py b/tests/test_e2e/test_ports.py
index 80365a01..9e46665a 100644
--- a/tests/test_e2e/test_ports.py
+++ b/tests/test_e2e/test_ports.py
@@ -2,7 +2,7 @@
 from ops import CharmBase, Framework, StartEvent, StopEvent
 
 from scenario import Context, State
-from scenario.state import StateValidationError, TCPPort, UDPPort, _Port
+from scenario.state import Port, StateValidationError, TCPPort, UDPPort
 
 
 class MyCharm(CharmBase):
@@ -42,7 +42,7 @@ def test_close_port(ctx):
 
 def test_port_no_arguments():
     with pytest.raises(RuntimeError):
-        _Port()
+        Port()
 
 
 @pytest.mark.parametrize("klass", (TCPPort, UDPPort))
diff --git a/tests/test_e2e/test_relations.py b/tests/test_e2e/test_relations.py
index 9ba0ed61..44433e21 100644
--- a/tests/test_e2e/test_relations.py
+++ b/tests/test_e2e/test_relations.py
@@ -17,10 +17,10 @@
     DEFAULT_JUJU_DATABAG,
     PeerRelation,
     Relation,
+    RelationBase,
     State,
     StateValidationError,
     SubordinateRelation,
-    _RelationBase,
     next_relation_id,
 )
 from tests.helpers import trigger
@@ -399,7 +399,7 @@ def post_event(charm: CharmBase):
 
 def test_cannot_instantiate_relationbase():
     with pytest.raises(RuntimeError):
-        _RelationBase("")
+        RelationBase("")
 
 
 def test_relation_ids():
@@ -417,7 +417,7 @@ def test_broken_relation_not_in_model_relations(mycharm):
     ctx = Context(
         mycharm, meta={"name": "local", "requires": {"foo": {"interface": "foo"}}}
     )
-    with ctx.manager(ctx.on.relation_broken(rel), state=State(relations={rel})) as mgr:
+    with ctx(ctx.on.relation_broken(rel), state=State(relations={rel})) as mgr:
         charm = mgr.charm
 
         assert charm.model.get_relation("foo") is None
diff --git a/tests/test_e2e/test_resource.py b/tests/test_e2e/test_resource.py
index c4237ea6..aebe8a0d 100644
--- a/tests/test_e2e/test_resource.py
+++ b/tests/test_e2e/test_resource.py
@@ -25,7 +25,7 @@ def test_get_resource():
     )
     resource1 = Resource(name="foo", path=pathlib.Path("/tmp/foo"))
     resource2 = Resource(name="bar", path=pathlib.Path("~/bar"))
-    with ctx.manager(
+    with ctx(
         ctx.on.update_status(), state=State(resources={resource1, resource2})
     ) as mgr:
         assert mgr.charm.model.resources.fetch("foo") == resource1.path
diff --git a/tests/test_e2e/test_secrets.py b/tests/test_e2e/test_secrets.py
index 5983c7df..a0249dad 100644
--- a/tests/test_e2e/test_secrets.py
+++ b/tests/test_e2e/test_secrets.py
@@ -28,7 +28,7 @@ def _on_event(self, event):
 
 def test_get_secret_no_secret(mycharm):
     ctx = Context(mycharm, meta={"name": "local"})
-    with ctx.manager(ctx.on.update_status(), State()) as mgr:
+    with ctx(ctx.on.update_status(), State()) as mgr:
         with pytest.raises(SecretNotFoundError):
             assert mgr.charm.model.get_secret(id="foo")
         with pytest.raises(SecretNotFoundError):
@@ -39,7 +39,7 @@ def test_get_secret_no_secret(mycharm):
 def test_get_secret(mycharm, owner):
     ctx = Context(mycharm, meta={"name": "local"})
     secret = Secret({"a": "b"}, owner=owner)
-    with ctx.manager(
+    with ctx(
         state=State(secrets={secret}),
         event=ctx.on.update_status(),
     ) as mgr:
@@ -54,7 +54,7 @@ def test_get_secret_get_refresh(mycharm, owner):
         latest_content={"a": "c"},
         owner=owner,
     )
-    with ctx.manager(
+    with ctx(
         ctx.on.update_status(),
         State(secrets={secret}),
     ) as mgr:
@@ -71,7 +71,7 @@ def test_get_secret_nonowner_peek_update(mycharm, app):
         tracked_content={"a": "b"},
         latest_content={"a": "c"},
     )
-    with ctx.manager(
+    with ctx(
         ctx.on.update_status(),
         State(
             leader=app,
@@ -98,7 +98,7 @@ def test_get_secret_owner_peek_update(mycharm, owner):
         latest_content={"a": "c"},
         owner=owner,
     )
-    with ctx.manager(
+    with ctx(
         ctx.on.update_status(),
         State(
             secrets={secret},
@@ -150,7 +150,7 @@ def test_consumer_events_failures(mycharm, evt_suffix, revision):
 @pytest.mark.parametrize("app", (True, False))
 def test_add(mycharm, app):
     ctx = Context(mycharm, meta={"name": "local"})
-    with ctx.manager(
+    with ctx(
         ctx.on.update_status(),
         State(leader=app),
     ) as mgr:
@@ -159,9 +159,10 @@ def test_add(mycharm, app):
             charm.app.add_secret({"foo": "bar"}, label="mylabel")
         else:
             charm.unit.add_secret({"foo": "bar"}, label="mylabel")
+        output = mgr.run()
 
-    assert mgr.output.secrets
-    secret = mgr.output.get_secret(label="mylabel")
+    assert output.secrets
+    secret = output.get_secret(label="mylabel")
     assert secret.latest_content == secret.tracked_content == {"foo": "bar"}
     assert secret.label == "mylabel"
 
@@ -171,7 +172,7 @@ def test_set_legacy_behaviour(mycharm):
     # ref: https://bugs.launchpad.net/juju/+bug/2037120
     ctx = Context(mycharm, meta={"name": "local"}, juju_version="3.1.6")
     rev1, rev2 = {"foo": "bar"}, {"foo": "baz", "qux": "roz"}
-    with ctx.manager(
+    with ctx(
         ctx.on.update_status(),
         State(),
     ) as mgr:
@@ -207,7 +208,7 @@ def test_set_legacy_behaviour(mycharm):
 def test_set(mycharm):
     ctx = Context(mycharm, meta={"name": "local"})
     rev1, rev2 = {"foo": "bar"}, {"foo": "baz", "qux": "roz"}
-    with ctx.manager(
+    with ctx(
         ctx.on.update_status(),
         State(),
     ) as mgr:
@@ -239,7 +240,7 @@ def test_set(mycharm):
 def test_set_juju33(mycharm):
     ctx = Context(mycharm, meta={"name": "local"}, juju_version="3.3.1")
     rev1, rev2 = {"foo": "bar"}, {"foo": "baz", "qux": "roz"}
-    with ctx.manager(
+    with ctx(
         ctx.on.update_status(),
         State(),
     ) as mgr:
@@ -271,7 +272,7 @@ def test_meta(mycharm, app):
         description="foobarbaz",
         rotate=SecretRotate.HOURLY,
     )
-    with ctx.manager(
+    with ctx(
         ctx.on.update_status(),
         State(
             leader=True,
@@ -308,7 +309,7 @@ def test_secret_permission_model(mycharm, leader, owner):
         rotate=SecretRotate.HOURLY,
     )
     secret_id = secret.id
-    with ctx.manager(
+    with ctx(
         ctx.on.update_status(),
         State(
             leader=leader,
@@ -355,7 +356,7 @@ def test_grant(mycharm, app):
         description="foobarbaz",
         rotate=SecretRotate.HOURLY,
     )
-    with ctx.manager(
+    with ctx(
         ctx.on.update_status(),
         State(
             relations=[Relation("foo", "remote")],
@@ -369,7 +370,8 @@ def test_grant(mycharm, app):
             secret.grant(relation=foo)
         else:
             secret.grant(relation=foo, unit=foo.units.pop())
-    vals = list(mgr.output.get_secret(label="mylabel").remote_grants.values())
+        output = mgr.run()
+    vals = list(output.get_secret(label="mylabel").remote_grants.values())
     assert vals == [{"remote"}] if app else [{"remote/0"}]
 
 
@@ -382,7 +384,7 @@ def test_update_metadata(mycharm):
         owner="unit",
         label="mylabel",
     )
-    with ctx.manager(
+    with ctx(
         ctx.on.update_status(),
         State(
             secrets={secret},
@@ -395,8 +397,9 @@ def test_update_metadata(mycharm):
             expire=exp,
             rotate=SecretRotate.DAILY,
         )
+        output = mgr.run()
 
-    secret_out = mgr.output.get_secret(label="babbuccia")
+    secret_out = output.get_secret(label="babbuccia")
     assert secret_out.label == "babbuccia"
     assert secret_out.rotate == SecretRotate.DAILY
     assert secret_out.description == "blu"
@@ -470,32 +473,35 @@ class GrantingCharm(CharmBase):
         },
     )
 
-    with ctx.manager(ctx.on.start(), state) as mgr:
+    with ctx(ctx.on.start(), state) as mgr:
         charm = mgr.charm
         secret = charm.app.add_secret({"foo": "bar"}, label="mylabel")
         bar_relation = charm.model.relations["bar"][0]
 
         secret.grant(bar_relation)
+        output = mgr.run()
 
-    assert mgr.output.secrets
-    scenario_secret = mgr.output.get_secret(label="mylabel")
+    assert output.secrets
+    scenario_secret = output.get_secret(label="mylabel")
     assert relation_remote_app in scenario_secret.remote_grants[relation_id]
 
-    with ctx.manager(ctx.on.start(), mgr.output) as mgr:
+    with ctx(ctx.on.start(), output) as mgr:
         charm: GrantingCharm = mgr.charm
         secret = charm.model.get_secret(label="mylabel")
         secret.revoke(bar_relation)
+        output = mgr.run()
 
-    scenario_secret = mgr.output.get_secret(label="mylabel")
+    scenario_secret = output.get_secret(label="mylabel")
     assert scenario_secret.remote_grants == {}
 
-    with ctx.manager(ctx.on.start(), mgr.output) as mgr:
+    with ctx(ctx.on.start(), output) as mgr:
         charm: GrantingCharm = mgr.charm
         secret = charm.model.get_secret(label="mylabel")
         secret.remove_all_revisions()
+        output = mgr.run()
 
     with pytest.raises(KeyError):
-        mgr.output.get_secret(label="mylabel")
+        output.get_secret(label="mylabel")
 
 
 def test_secret_removed_event():
diff --git a/tests/test_e2e/test_storage.py b/tests/test_e2e/test_storage.py
index 3e6912fb..34785aa1 100644
--- a/tests/test_e2e/test_storage.py
+++ b/tests/test_e2e/test_storage.py
@@ -23,13 +23,13 @@ def no_storage_ctx():
 
 
 def test_storage_get_null(no_storage_ctx):
-    with no_storage_ctx.manager(no_storage_ctx.on.update_status(), State()) as mgr:
+    with no_storage_ctx(no_storage_ctx.on.update_status(), State()) as mgr:
         storages = mgr.charm.model.storages
         assert not len(storages)
 
 
 def test_storage_get_unknown_name(storage_ctx):
-    with storage_ctx.manager(storage_ctx.on.update_status(), State()) as mgr:
+    with storage_ctx(storage_ctx.on.update_status(), State()) as mgr:
         storages = mgr.charm.model.storages
         # not in metadata
         with pytest.raises(KeyError):
@@ -37,7 +37,7 @@ def test_storage_get_unknown_name(storage_ctx):
 
 
 def test_storage_request_unknown_name(storage_ctx):
-    with storage_ctx.manager(storage_ctx.on.update_status(), State()) as mgr:
+    with storage_ctx(storage_ctx.on.update_status(), State()) as mgr:
         storages = mgr.charm.model.storages
         # not in metadata
         with pytest.raises(ModelError):
@@ -45,7 +45,7 @@ def test_storage_request_unknown_name(storage_ctx):
 
 
 def test_storage_get_some(storage_ctx):
-    with storage_ctx.manager(storage_ctx.on.update_status(), State()) as mgr:
+    with storage_ctx(storage_ctx.on.update_status(), State()) as mgr:
         storages = mgr.charm.model.storages
         # known but none attached
         assert storages["foo"] == []
@@ -53,7 +53,7 @@ def test_storage_get_some(storage_ctx):
 
 @pytest.mark.parametrize("n", (1, 3, 5))
 def test_storage_add(storage_ctx, n):
-    with storage_ctx.manager(storage_ctx.on.update_status(), State()) as mgr:
+    with storage_ctx(storage_ctx.on.update_status(), State()) as mgr:
         storages = mgr.charm.model.storages
         storages.request("foo", n)
 
@@ -65,9 +65,7 @@ def test_storage_usage(storage_ctx):
     # setup storage with some content
     (storage.get_filesystem(storage_ctx) / "myfile.txt").write_text("helloworld")
 
-    with storage_ctx.manager(
-        storage_ctx.on.update_status(), State(storages={storage})
-    ) as mgr:
+    with storage_ctx(storage_ctx.on.update_status(), State(storages={storage})) as mgr:
         foo = mgr.charm.model.storages["foo"][0]
         loc = foo.location
         path = loc / "myfile.txt"
diff --git a/tests/test_e2e/test_vroot.py b/tests/test_e2e/test_vroot.py
index c8702611..fa8a3d41 100644
--- a/tests/test_e2e/test_vroot.py
+++ b/tests/test_e2e/test_vroot.py
@@ -56,7 +56,7 @@ def test_charm_virtual_root_cleanup_if_exists(charm_virtual_root):
     meta_file.write_text(raw_ori_meta)
 
     ctx = Context(MyCharm, meta=MyCharm.META, charm_root=charm_virtual_root)
-    with ctx.manager(
+    with ctx(
         ctx.on.start(),
         State(),
     ) as mgr:
@@ -79,7 +79,7 @@ def test_charm_virtual_root_cleanup_if_not_exists(charm_virtual_root):
     assert not meta_file.exists()
 
     ctx = Context(MyCharm, meta=MyCharm.META, charm_root=charm_virtual_root)
-    with ctx.manager(
+    with ctx(
         ctx.on.start(),
         State(),
     ) as mgr: