Merge pull request #78 from canonical/collect-status-support

Adds support for `collect-status`

README.md

@@ -238,7 +238,7 @@ assert ctx.workload_version_history == ['1', '1.2', '1.5']
 
 If your charm deals with deferred events, custom events, and charm libs that in turn emit their own custom events, it
 can be hard to examine the resulting control flow. In these situations it can be useful to verify that, as a result of a
-given juju event triggering (say, 'start'), a specific chain of deferred and custom events is emitted on the charm. The
+given Juju event triggering (say, 'start'), a specific chain of events is emitted on the charm. The
 resulting state, black-box as it is, gives little insight into how exactly it was obtained.
 
 ```python
@@ -254,6 +254,33 @@ def test_foo():
     assert isinstance(ctx.emitted_events[0], StartEvent)
 ```
 
+You can configure what events will be captured by passing the following arguments to `Context`:
+-  `capture_deferred_events`: If you want to include re-emitted deferred events.
+-  `capture_framework_events`: If you want to include framework events (`pre-commit`, `commit`, and `collect-status`). 
+
+For example:
+```python
+from scenario import Context, Event, State
+
+def test_emitted_full():
+    ctx = Context(
+        MyCharm,
+        capture_deferred_events=True,
+        capture_framework_events=True,
+    )
+    ctx.run("start", State(deferred=[Event("update-status").deferred(MyCharm._foo)]))
+
+    assert len(ctx.emitted_events) == 5
+    assert [e.handle.kind for e in ctx.emitted_events] == [
+        "update_status",
+        "start",
+        "collect_unit_status",
+        "pre_commit",
+        "commit",
+    ]
+```
+
+
 ### Low-level access: using directly `capture_events`
 
 If you need more control over what events are captured (or you're not into pytest), you can use directly the context
@@ -299,7 +326,7 @@ Configuration:
 - By default, **framework events** (`PreCommit`, `Commit`) are not considered for inclusion in the output list even if
   they match the instance check. You can toggle that by passing: `capture_events(include_framework=True)`.
 - By default, **deferred events** are included in the listing if they match the instance check. You can toggle that by
-  passing: `capture_events(include_deferred=True)`.
+  passing: `capture_events(include_deferred=False)`.
 
 ## Relations
 

pyproject.toml

@@ -8,7 +8,7 @@ build-backend = "setuptools.build_meta"
 [project]
 name = "ops-scenario"
 
-version = "5.5"
+version = "5.6"
 
 authors = [
     { name = "Pietro Pasotti", email = "[email protected]" }

scenario/capture_events.py

@@ -6,6 +6,7 @@
 from contextlib import contextmanager
 from typing import ContextManager, List, Type, TypeVar
 
+from ops import CollectStatusEvent
 from ops.framework import (
     CommitEvent,
     EventBase,
@@ -49,7 +50,10 @@ def capture_events(
     _real_reemit = Framework.reemit
 
     def _wrapped_emit(self, evt):
-        if not include_framework and isinstance(evt, (PreCommitEvent, CommitEvent)):
+        if not include_framework and isinstance(
+            evt,
+            (PreCommitEvent, CommitEvent, CollectStatusEvent),
+        ):
             return _real_emit(self, evt)
 
         if isinstance(evt, allowed_types):

scenario/context.py

@@ -167,6 +167,8 @@ def __init__(
         config: Optional[Dict[str, Any]] = None,
         charm_root: "PathLike" = None,
         juju_version: str = "3.0",
+        capture_deferred_events: bool = False,
+        capture_framework_events: bool = False,
     ):
         """Represents a simulated charm's execution context.
 
@@ -258,6 +260,10 @@ def __init__(
         self.juju_version = juju_version
         self._tmp = tempfile.TemporaryDirectory()
 
+        # config for what events to be captured in emitted_events.
+        self.capture_deferred_events = capture_deferred_events
+        self.capture_framework_events = capture_framework_events
+
         # streaming side effects from running an event
         self.juju_log: List["JujuLogLine"] = []
         self.app_status_history: List["_EntityStatus"] = []

scenario/ops_main_mock.py

@@ -195,7 +195,11 @@ def commit(self):
         if not self._has_emitted:
             raise RuntimeError("should .emit() before you .commit()")
 
+        # emit collect-status events
+        ops.charm._evaluate_status(self.charm)
+
         self._has_committed = True
+
         try:
             self.framework.commit()
         finally:

scenario/runtime.py

@@ -376,11 +376,14 @@ def _close_storage(self, state: "State", temporary_charm_root: Path):
         return state.replace(deferred=deferred, stored_state=stored_state)
 
     @contextmanager
-    def _exec_ctx(self) -> ContextManager[Tuple[Path, List[EventBase]]]:
+    def _exec_ctx(self, ctx: "Context") -> ContextManager[Tuple[Path, List[EventBase]]]:
         """python 3.8 compatibility shim"""
         with self._virtual_charm_root() as temporary_charm_root:
             # todo allow customizing capture_events
-            with capture_events() as captured:
+            with capture_events(
+                include_deferred=ctx.capture_deferred_events,
+                include_framework=ctx.capture_framework_events,
+            ) as captured:
                 yield (temporary_charm_root, captured)
 
     @contextmanager
@@ -412,7 +415,7 @@ def exec(
         output_state = state.copy()
 
         logger.info(" - generating virtual charm root")
-        with self._exec_ctx() as (temporary_charm_root, captured):
+        with self._exec_ctx(context) as (temporary_charm_root, captured):
             logger.info(" - initializing storage")
             self._initialize_storage(state, temporary_charm_root)
 

scenario/state.py

@@ -62,6 +62,10 @@
     "leader_elected",
     "leader_settings_changed",
     "collect_metrics",
+}
+FRAMEWORK_EVENTS = {
+    "pre_commit",
+    "commit",
     "collect_app_status",
     "collect_unit_status",
 }
@@ -998,6 +1002,7 @@ def name(self):
 
 
 class _EventType(str, Enum):
+    framework = "framework"
     builtin = "builtin"
     relation = "relation"
     action = "action"
@@ -1038,6 +1043,9 @@ def _get_suffix_and_type(s: str):
         if s in SECRET_EVENTS:
             return s, _EventType.secret
 
+        if s in FRAMEWORK_EVENTS:
+            return s, _EventType.framework
+
         # Whether the event name indicates that this is a storage event.
         for suffix in STORAGE_EVENTS_SUFFIX:
             if s.endswith(suffix):

tests/test_e2e/test_actions.py

@@ -97,7 +97,9 @@ def test_cannot_run_action_event(mycharm):
 
 @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: ActionEvent):
+    def handle_evt(charm: CharmBase, evt):
+        if not isinstance(evt, ActionEvent):
+            return
         evt.set_results(res_value)
         evt.log("foo")
         evt.log("bar")
@@ -116,6 +118,9 @@ def handle_evt(charm: CharmBase, evt: ActionEvent):
 @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):
+        if not isinstance(evt, ActionEvent):
+            return
+
         evt.set_results({"my-res": res_value})
         evt.log("log1")
         evt.log("log2")

tests/test_e2e/test_builtin_scenes.py

@@ -1,5 +1,5 @@
 import pytest
-from ops.charm import CharmBase
+from ops.charm import CharmBase, CollectStatusEvent
 from ops.framework import Framework
 
 from scenario.sequences import check_builtin_sequences
@@ -27,6 +27,9 @@ def __init__(self, framework: Framework):
                 self.framework.observe(evt, self._on_event)
 
         def _on_event(self, event):
+            if isinstance(event, CollectStatusEvent):
+                return
+
             global CHARM_CALLED
             CHARM_CALLED += 1
 

tests/test_e2e/test_deferred.py

@@ -3,6 +3,7 @@
 import pytest
 from ops.charm import (
     CharmBase,
+    CollectStatusEvent,
     RelationChangedEvent,
     StartEvent,
     UpdateStatusEvent,
@@ -65,8 +66,8 @@ def test_deferred_evt_emitted(mycharm):
     assert out.deferred[1].name == "update_status"
 
     # we saw start and update-status.
-    assert len(mycharm.captured) == 2
-    upstat, start = mycharm.captured
+    assert len(mycharm.captured) == 3
+    upstat, start, _ = mycharm.captured
     assert isinstance(upstat, UpdateStatusEvent)
     assert isinstance(start, StartEvent)
 
@@ -123,8 +124,8 @@ def test_deferred_relation_event(mycharm):
     assert out.deferred[1].name == "start"
 
     # we saw start and relation-changed.
-    assert len(mycharm.captured) == 2
-    upstat, start = mycharm.captured
+    assert len(mycharm.captured) == 3
+    upstat, start, _ = mycharm.captured
     assert isinstance(upstat, RelationChangedEvent)
     assert isinstance(start, StartEvent)
 
@@ -156,10 +157,11 @@ def test_deferred_relation_event_from_relation(mycharm):
     assert out.deferred[1].name == "start"
 
     # we saw start and foo_relation_changed.
-    assert len(mycharm.captured) == 2
-    upstat, start = mycharm.captured
+    assert len(mycharm.captured) == 3
+    upstat, start, collect_status = mycharm.captured
     assert isinstance(upstat, RelationChangedEvent)
     assert isinstance(start, StartEvent)
+    assert isinstance(collect_status, CollectStatusEvent)
 
 
 def test_deferred_workload_event(mycharm):
@@ -183,14 +185,15 @@ def test_deferred_workload_event(mycharm):
     assert out.deferred[1].name == "start"
 
     # we saw start and foo_pebble_ready.
-    assert len(mycharm.captured) == 2
-    upstat, start = mycharm.captured
+    assert len(mycharm.captured) == 3
+    upstat, start, collect_status = mycharm.captured
     assert isinstance(upstat, WorkloadEvent)
     assert isinstance(start, StartEvent)
+    assert isinstance(collect_status, CollectStatusEvent)
 
 
 def test_defer_reemit_lifecycle_event(mycharm):
-    ctx = Context(mycharm, meta=mycharm.META)
+    ctx = Context(mycharm, meta=mycharm.META, capture_deferred_events=True)
 
     mycharm.defer_next = 1
     state_1 = ctx.run("update-status", State())
@@ -208,7 +211,7 @@ def test_defer_reemit_lifecycle_event(mycharm):
 
 
 def test_defer_reemit_relation_event(mycharm):
-    ctx = Context(mycharm, meta=mycharm.META)
+    ctx = Context(mycharm, meta=mycharm.META, capture_deferred_events=True)
 
     rel = Relation("foo")
     mycharm.defer_next = 1

tests/test_e2e/test_event.py

@@ -1,7 +1,9 @@
+import ops
 import pytest
-from ops import CharmBase
+from ops import CharmBase, StartEvent, UpdateStatusEvent
 
-from scenario.state import Event, _CharmSpec, _EventType
+from scenario import Context
+from scenario.state import Event, State, _CharmSpec, _EventType
 
 
 @pytest.mark.parametrize(
@@ -16,6 +18,10 @@
         ("foo_pebble_ready", _EventType.workload),
         ("foo_bar_baz_pebble_ready", _EventType.workload),
         ("secret_removed", _EventType.secret),
+        ("pre_commit", _EventType.framework),
+        ("commit", _EventType.framework),
+        ("collect_unit_status", _EventType.framework),
+        ("collect_app_status", _EventType.framework),
         ("foo", _EventType.custom),
         ("kaboozle_bar_baz", _EventType.custom),
     ),
@@ -48,3 +54,43 @@ class MyCharm(CharmBase):
         },
     )
     assert event._is_builtin_event(spec) is (expected_type is not _EventType.custom)
+
+
+def test_emitted_framework():
+    class MyCharm(CharmBase):
+        META = {"name": "joop"}
+
+    ctx = Context(MyCharm, meta=MyCharm.META, capture_framework_events=True)
+    ctx.run("update-status", State())
+    assert len(ctx.emitted_events) == 4
+    assert list(map(type, ctx.emitted_events)) == [
+        ops.UpdateStatusEvent,
+        ops.CollectStatusEvent,
+        ops.PreCommitEvent,
+        ops.CommitEvent,
+    ]
+
+
+def test_emitted_deferred():
+    class MyCharm(CharmBase):
+        META = {"name": "joop"}
+
+        def _foo(self, e):
+            pass
+
+    ctx = Context(
+        MyCharm,
+        meta=MyCharm.META,
+        capture_deferred_events=True,
+        capture_framework_events=True,
+    )
+    ctx.run("start", State(deferred=[Event("update-status").deferred(MyCharm._foo)]))
+
+    assert len(ctx.emitted_events) == 5
+    assert [e.handle.kind for e in ctx.emitted_events] == [
+        "update_status",
+        "start",
+        "collect_unit_status",
+        "pre_commit",
+        "commit",
+    ]

tests/test_e2e/test_juju_log.py

@@ -1,7 +1,7 @@
 import logging
 
 import pytest
-from ops.charm import CharmBase
+from ops.charm import CharmBase, CollectStatusEvent
 
 from scenario import Context
 from scenario.state import JujuLogLine, State
@@ -21,6 +21,8 @@ def __init__(self, framework):
                 self.framework.observe(evt, self._on_event)
 
         def _on_event(self, event):
+            if isinstance(event, CollectStatusEvent):
+                return
             print("foo!")
             logger.warning("bar!")
 

tests/test_e2e/test_manager.py

@@ -2,7 +2,7 @@
 
 import pytest
 from ops import ActiveStatus
-from ops.charm import CharmBase
+from ops.charm import CharmBase, CollectStatusEvent
 
 from scenario import Action, Context, State
 from scenario.context import ActionOutput, AlreadyEmittedError, _EventManager
@@ -20,6 +20,9 @@ def __init__(self, framework):
                 self.framework.observe(evt, self._on_event)
 
         def _on_event(self, e):
+            if isinstance(e, CollectStatusEvent):
+                return
+
             print("event!")
             self.unit.status = ActiveStatus(e.handle.kind)