feat: add support for Pebble check events (#151)

* Add support for Pebble checks.

Also update the support for Pebble notices to be aligned with the 7.x approach.

README.md

@@ -691,7 +691,25 @@ notices = [
     scenario.Notice(key="example.com/c"),
 ]
 container = scenario.Container("my-container", notices=notices)
-ctx.run(container.get_notice("example.com/c").event, scenario.State(containers=[container]))
+state = scenario.State(containers={container})
+ctx.run(ctx.on.pebble_custom_notice(container=container, notice=notices[-1]), state)
+```
+
+### Pebble Checks
+
+A Pebble plan can contain checks, and when those checks exceed the configured
+failure threshold, or start succeeding again after, Juju will emit a
+pebble-check-failed or pebble-check-recovered event. In order to simulate these
+events, you need to add a `CheckInfo` to the container. Note that the status of the
+check doesn't have to match the event being generated: by the time that Juju
+sends a pebble-check-failed event the check might have started passing again.
+
+```python
+ctx = scenario.Context(MyCharm, meta={"name": "foo", "containers": {"my-container": {}}})
+check_info = scenario.CheckInfo("http-check", failures=7, status=ops.pebble.CheckStatus.DOWN)
+container = scenario.Container("my-container", check_infos={check_info})
+state = scenario.State(containers={container})
+ctx.run(ctx.on.pebble_check_failed(info=check_info, container=container), state=state)
 ```
 
 ## Storage

pyproject.toml

@@ -18,7 +18,7 @@ license.text = "Apache-2.0"
 keywords = ["juju", "test"]
 
 dependencies = [
-    "ops>=2.12",
+    "ops>=2.15",
     "PyYAML>=6.0.1",
 ]
 readme = "README.md"

scenario/__init__.py

@@ -7,6 +7,7 @@
     Address,
     BindAddress,
     BlockedStatus,
+    CheckInfo,
     CloudCredential,
     CloudSpec,
     Container,
@@ -37,6 +38,7 @@
 
 __all__ = [
     "ActionOutput",
+    "CheckInfo",
     "CloudCredential",
     "CloudSpec",
     "Context",

scenario/consistency_checker.py

@@ -563,6 +563,9 @@ def check_containers_consistency(
     meta_containers = list(map(normalize_name, meta.get("containers", {})))
     state_containers = [normalize_name(c.name) for c in state.containers]
     all_notices = {notice.id for c in state.containers for notice in c.notices}
+    all_checks = {
+        (c.name, check.name) for c in state.containers for check in c.check_infos
+    }
     errors = []
 
     # it's fine if you have containers in meta that are not in state.containers (yet), but it's
@@ -587,6 +590,14 @@ def check_containers_consistency(
                 f"the event being processed concerns notice {event.notice!r}, but that "
                 "notice is not in any of the containers present in the state.",
             )
+        if (
+            event.check_info
+            and (evt_container_name, event.check_info.name) not in all_checks
+        ):
+            errors.append(
+                f"the event being processed concerns check {event.check_info.name}, but that "
+                "check is not the {evt_container_name} container.",
+            )
 
     # - a container in state.containers is not in meta.containers
     if diff := (set(state_containers).difference(set(meta_containers))):

scenario/context.py

@@ -12,8 +12,10 @@
 from scenario.logger import logger as scenario_logger
 from scenario.runtime import Runtime
 from scenario.state import (
+    CheckInfo,
     Container,
     MetadataNotFoundError,
+    Notice,
     Secret,
     Storage,
     _Action,
@@ -312,6 +314,30 @@ def storage_detaching(storage: Storage):
     def pebble_ready(container: Container):
         return _Event(f"{container.name}_pebble_ready", container=container)
 
+    @staticmethod
+    def pebble_custom_notice(container: Container, notice: Notice):
+        return _Event(
+            f"{container.name}_pebble_custom_notice",
+            container=container,
+            notice=notice,
+        )
+
+    @staticmethod
+    def pebble_check_failed(container: Container, info: CheckInfo):
+        return _Event(
+            f"{container.name}_pebble_check_failed",
+            container=container,
+            check_info=info,
+        )
+
+    @staticmethod
+    def pebble_check_recovered(container: Container, info: CheckInfo):
+        return _Event(
+            f"{container.name}_pebble_check_recovered",
+            container=container,
+            check_info=info,
+        )
+
     @staticmethod
     def action(
         name: str,

scenario/mocking.py

@@ -694,15 +694,18 @@ def __init__(
 
         self._root = container_root
 
-        # load any existing notices from the state
+        # load any existing notices and check information from the state
         self._notices: Dict[Tuple[str, str], pebble.Notice] = {}
+        self._check_infos: Dict[str, pebble.CheckInfo] = {}
         for container in state.containers:
             for notice in container.notices:
                 if hasattr(notice.type, "value"):
                     notice_type = cast(pebble.NoticeType, notice.type).value
                 else:
                     notice_type = str(notice.type)
                 self._notices[notice_type, notice.key] = notice._to_ops()
+            for check in container.check_infos:
+                self._check_infos[check.name] = check._to_ops()
 
     def get_plan(self) -> pebble.Plan:
         return self._container.plan

scenario/runtime.py

@@ -271,6 +271,9 @@ def _get_event_env(self, state: "State", event: "_Event", charm_root: Path):
                 },
             )
 
+        if check_info := event.check_info:
+            env["JUJU_PEBBLE_CHECK_NAME"] = check_info.name
+
         if storage := event.storage:
             env.update({"JUJU_STORAGE_ID": f"{storage.name}/{storage.index}"})
 

scenario/state.py

@@ -91,6 +91,8 @@
 }
 PEBBLE_READY_EVENT_SUFFIX = "_pebble_ready"
 PEBBLE_CUSTOM_NOTICE_EVENT_SUFFIX = "_pebble_custom_notice"
+PEBBLE_CHECK_FAILED_EVENT_SUFFIX = "_pebble_check_failed"
+PEBBLE_CHECK_RECOVERED_EVENT_SUFFIX = "_pebble_check_recovered"
 RELATION_EVENTS_SUFFIX = {
     "_relation_changed",
     "_relation_broken",
@@ -770,18 +772,37 @@ def _to_ops(self) -> pebble.Notice:
 
 
 @dataclasses.dataclass(frozen=True)
-class _BoundNotice(_max_posargs(0)):
-    notice: Notice
-    container: "Container"
+class CheckInfo(_max_posargs(1)):
+    name: str
+    """Name of the check."""
 
-    @property
-    def event(self):
-        """Sugar to generate a <container's name>-pebble-custom-notice event for this notice."""
-        suffix = PEBBLE_CUSTOM_NOTICE_EVENT_SUFFIX
-        return _Event(
-            path=normalize_name(self.container.name) + suffix,
-            container=self.container,
-            notice=self.notice,
+    level: Optional[pebble.CheckLevel] = None
+    """Level of the check."""
+
+    status: pebble.CheckStatus = pebble.CheckStatus.UP
+    """Status of the check.
+
+    CheckStatus.UP means the check is healthy (the number of failures is less
+    than the threshold), CheckStatus.DOWN means the check is unhealthy
+    (the number of failures has reached the threshold).
+    """
+
+    failures: int = 0
+    """Number of failures since the check last succeeded."""
+
+    threshold: int = 3
+    """Failure threshold.
+
+    This is how many consecutive failures for the check to be considered “down”.
+    """
+
+    def _to_ops(self) -> pebble.CheckInfo:
+        return pebble.CheckInfo(
+            name=self.name,
+            level=self.level,
+            status=self.status,
+            failures=self.failures,
+            threshold=self.threshold,
         )
 
 
@@ -862,6 +883,8 @@ class Container(_max_posargs(1)):
 
     notices: List[Notice] = dataclasses.field(default_factory=list)
 
+    check_infos: FrozenSet[CheckInfo] = frozenset()
+
     def __hash__(self) -> int:
         return hash(self.name)
 
@@ -927,23 +950,6 @@ def get_filesystem(self, ctx: "Context") -> Path:
         """
         return ctx._get_container_root(self.name)
 
-    def get_notice(
-        self,
-        key: str,
-        notice_type: pebble.NoticeType = pebble.NoticeType.CUSTOM,
-    ) -> _BoundNotice:
-        """Get a Pebble notice by key and type.
-
-        Raises:
-            KeyError: if the notice is not found.
-        """
-        for notice in self.notices:
-            if notice.key == key and notice.type == notice_type:
-                return _BoundNotice(notice=notice, container=self)
-        raise KeyError(
-            f"{self.name} does not have a notice with key {key} and type {notice_type}",
-        )
-
 
 _RawStatusLiteral = Literal[
     "waiting",
@@ -1631,6 +1637,10 @@ def _get_suffix_and_type(s: str) -> Tuple[str, _EventType]:
             return PEBBLE_READY_EVENT_SUFFIX, _EventType.workload
         if s.endswith(PEBBLE_CUSTOM_NOTICE_EVENT_SUFFIX):
             return PEBBLE_CUSTOM_NOTICE_EVENT_SUFFIX, _EventType.workload
+        if s.endswith(PEBBLE_CHECK_FAILED_EVENT_SUFFIX):
+            return PEBBLE_CHECK_FAILED_EVENT_SUFFIX, _EventType.workload
+        if s.endswith(PEBBLE_CHECK_RECOVERED_EVENT_SUFFIX):
+            return PEBBLE_CHECK_RECOVERED_EVENT_SUFFIX, _EventType.workload
 
         if s in BUILTIN_EVENTS:
             return "", _EventType.builtin
@@ -1670,6 +1680,9 @@ class Event:
     notice: Optional[Notice] = None
     """If this is a Pebble notice event, the notice it refers to."""
 
+    check_info: Optional[CheckInfo] = None
+    """If this is a Pebble check event, the check info it provides."""
+
     action: Optional["_Action"] = None
     """If this is an action event, the :class:`Action` it refers to."""
 
@@ -1787,6 +1800,8 @@ def deferred(self, handler: Callable, event_id: int = 1) -> DeferredEvent:
                         "notice_type": notice_type,
                     },
                 )
+            elif self.check_info:
+                snapshot_data["check_name"] = self.check_info.name
 
         elif self._is_relation_event:
             # this is a RelationEvent.
@@ -1860,8 +1875,15 @@ def deferred(
     relation: Optional["Relation"] = None,
     container: Optional["Container"] = None,
     notice: Optional["Notice"] = None,
+    check_info: Optional["CheckInfo"] = None,
 ):
     """Construct a DeferredEvent from an Event or an event name."""
     if isinstance(event, str):
-        event = _Event(event, relation=relation, container=container, notice=notice)
+        event = _Event(
+            event,
+            relation=relation,
+            container=container,
+            notice=notice,
+            check_info=check_info,
+        )
     return event.deferred(handler=handler, event_id=event_id)

tests/test_consistency_checker.py

@@ -8,6 +8,7 @@
 from scenario.runtime import InconsistentScenarioError
 from scenario.state import (
     RELATION_EVENTS_SUFFIX,
+    CheckInfo,
     CloudCredential,
     CloudSpec,
     Container,
@@ -85,6 +86,46 @@ def test_workload_event_without_container():
         _Event("foo-pebble-custom-notice", container=Container("foo"), notice=notice),
         _CharmSpec(MyCharm, {"containers": {"foo": {}}}),
     )
+    check = CheckInfo("http-check")
+    assert_consistent(
+        State(containers={Container("foo", check_infos={check})}),
+        _Event("foo-pebble-check-failed", container=Container("foo"), check_info=check),
+        _CharmSpec(MyCharm, {"containers": {"foo": {}}}),
+    )
+    assert_inconsistent(
+        State(containers={Container("foo")}),
+        _Event("foo-pebble-check-failed", container=Container("foo"), check_info=check),
+        _CharmSpec(MyCharm, {"containers": {"foo": {}}}),
+    )
+    assert_consistent(
+        State(containers={Container("foo", check_infos={check})}),
+        _Event(
+            "foo-pebble-check-recovered", container=Container("foo"), check_info=check
+        ),
+        _CharmSpec(MyCharm, {"containers": {"foo": {}}}),
+    )
+    assert_inconsistent(
+        State(containers={Container("foo")}),
+        _Event(
+            "foo-pebble-check-recovered", container=Container("foo"), check_info=check
+        ),
+        _CharmSpec(MyCharm, {"containers": {"foo": {}}}),
+    )
+    # Ensure the check is in the correct container.
+    assert_inconsistent(
+        State(containers={Container("foo", check_infos={check}), Container("bar")}),
+        _Event(
+            "foo-pebble-check-recovered", container=Container("bar"), check_info=check
+        ),
+        _CharmSpec(MyCharm, {"containers": {"foo": {}, "bar": {}}}),
+    )
+    assert_inconsistent(
+        State(containers={Container("foo", check_infos={check}), Container("bar")}),
+        _Event(
+            "bar-pebble-check-recovered", container=Container("bar"), check_info=check
+        ),
+        _CharmSpec(MyCharm, {"containers": {"foo": {}, "bar": {}}}),
+    )
 
 
 def test_container_meta_mismatch():

tests/test_e2e/test_deferred.py

@@ -102,7 +102,9 @@ def test_deferred_workload_evt(mycharm):
 def test_deferred_notice_evt(mycharm):
     notice = Notice(key="example.com/bar")
     ctr = Container("foo", notices=[notice])
-    evt1 = ctr.get_notice("example.com/bar").event.deferred(handler=mycharm._on_event)
+    evt1 = _Event("foo_pebble_custom_notice", notice=notice, container=ctr).deferred(
+        handler=mycharm._on_event
+    )
     evt2 = deferred(
         event="foo_pebble_custom_notice",
         handler=mycharm._on_event,