diff --git a/scenario/context.py b/scenario/context.py index 5cf51635..7167d14e 100644 --- a/scenario/context.py +++ b/scenario/context.py @@ -599,7 +599,7 @@ def _record_status(self, state: "State", is_app: bool): else: self.unit_status_history.append(state.unit_status) - def __call__(self, event: Union["_Event", "Action"], state: "State"): + def __call__(self, event: Union["_Event", "_Action"], state: "State"): """Context manager to introspect live charm object before and after the event is emitted. Usage:: @@ -619,7 +619,7 @@ def __call__(self, event: Union["_Event", "Action"], state: "State"): event: the :class:`Event` or :class:`Action` that the charm will respond to. state: the :class:`State` instance to use when handling the Event. """ - if isinstance(event, Action) or event.action: + if isinstance(event, _Action) or event.action: return _ActionManager(self, event, state) return _EventManager(self, event, state) diff --git a/tests/test_context.py b/tests/test_context.py index f3ad4d16..851e73f4 100644 --- a/tests/test_context.py +++ b/tests/test_context.py @@ -4,7 +4,7 @@ from ops import CharmBase from scenario import ActionOutput, Context, State -from scenario.state import _Action, _Event, next_action_id +from scenario.state import _Event, next_action_id class MyCharm(CharmBase): @@ -65,7 +65,7 @@ def test_context_manager(): event.run() assert event.charm.meta.name == "foo" - with ctx(Action("act"), state) as event: + with ctx(ctx.on.action("act"), state) as event: event.run_action() assert event.charm.meta.name == "foo"