From 68cf2d7ce03865e4b23091ac8dbc3757aae4d4b7 Mon Sep 17 00:00:00 2001
From: Tony Meyer <tony.meyer@canonical.com>
Date: Fri, 11 Oct 2024 17:05:41 +1300
Subject: [PATCH] fix: raise ModelError on unknown/error status set in Scenario
 (#1417)

This copies across the fix from Harness to also be present in Scenario:
if `status_set` is passed anything other than the valid statuses, then
`ModelError` is raised (to better mimic what the Juju hook would
actually do).

Closes https://github.com/canonical/ops-scenario/issues/202 (rather than
copying it over to this repo).

---------

Co-authored-by: Ben Hoyt <benhoyt@gmail.com>
---
 testing/src/mocking.py                | 10 ++++--
 testing/tests/test_e2e/test_status.py | 46 +++++++++++++++++++++++++++
 2 files changed, 54 insertions(+), 2 deletions(-)

diff --git a/testing/src/mocking.py b/testing/src/mocking.py
index 534b4cabc..97bfe97f1 100644
--- a/testing/src/mocking.py
+++ b/testing/src/mocking.py
@@ -20,6 +20,7 @@
     Tuple,
     Union,
     cast,
+    get_args,
 )
 
 from ops import JujuVersion, pebble
@@ -35,6 +36,7 @@
     SecretRotate,
     _format_action_result_dict,
     _ModelBackend,
+    _SettableStatusName,
 )
 from ops.pebble import Client, ExecError
 
@@ -53,7 +55,6 @@
     _EntityStatus,
     _port_cls_by_protocol,
     _RawPortProtocolLiteral,
-    _RawStatusLiteral,
 )
 
 if TYPE_CHECKING:  # pragma: no cover
@@ -351,11 +352,16 @@ def application_version_set(self, version: str):
 
     def status_set(
         self,
-        status: _RawStatusLiteral,
+        status: _SettableStatusName,
         message: str = "",
         *,
         is_app: bool = False,
     ):
+        valid_names = get_args(_SettableStatusName)
+        if status not in valid_names:
+            raise ModelError(
+                f'ERROR invalid status "{status}", expected one of [{", ".join(valid_names)}]',
+            )
         self._context._record_status(self._state, is_app)
         status_obj = _EntityStatus.from_status_name(status, message)
         self._state._update_status(status_obj, is_app)
diff --git a/testing/tests/test_e2e/test_status.py b/testing/tests/test_e2e/test_status.py
index 2f98c6f8b..330819a3d 100644
--- a/testing/tests/test_e2e/test_status.py
+++ b/testing/tests/test_e2e/test_status.py
@@ -13,6 +13,7 @@
     UnknownStatus,
     WaitingStatus,
 )
+from scenario.errors import UncaughtCharmError
 from ..helpers import trigger
 
 
@@ -169,3 +170,48 @@ def test_status_comparison(status):
     assert isinstance(status, type(ops_status))
     # The repr of the scenario and ops classes should be identical.
     assert repr(status) == repr(ops_status)
+
+
+@pytest.mark.parametrize(
+    "status",
+    (
+        ActiveStatus("foo"),
+        WaitingStatus("bar"),
+        BlockedStatus("baz"),
+        MaintenanceStatus("qux"),
+    ),
+)
+def test_status_success(status: ops.StatusBase):
+    class MyCharm(CharmBase):
+        def __init__(self, framework: Framework):
+            super().__init__(framework)
+            framework.observe(self.on.update_status, self._on_update_status)
+
+        def _on_update_status(self, _):
+            self.unit.status = status
+
+    ctx = Context(MyCharm, meta={"name": "foo"})
+    ctx.run(ctx.on.update_status(), State())
+
+
+@pytest.mark.parametrize(
+    "status",
+    (
+        ErrorStatus("fiz"),
+        UnknownStatus(),
+    ),
+)
+def test_status_error(status: ops.StatusBase):
+    class MyCharm(CharmBase):
+        def __init__(self, framework: Framework):
+            super().__init__(framework)
+            framework.observe(self.on.update_status, self._on_update_status)
+
+        def _on_update_status(self, _):
+            self.unit.status = status
+
+    ctx = Context(MyCharm, meta={"name": "foo"})
+    with pytest.raises(UncaughtCharmError) as excinfo:
+        ctx.run(ctx.on.update_status(), State())
+    assert isinstance(excinfo.value.__cause__, ops.ModelError)
+    assert f'invalid status "{status.name}"' in str(excinfo.value.__cause__)