Repository: canonical/operator
Pull request #1354 (2024-09-03 05:47:59)

chore: Document and validate settable status values in _ModelBackend.set_status
Document the status values that are valid arguments for `_ModelBackend.status_set`, as requested in #1343, via type annotation. The `status` argument was previously type hinted only as `str`. However, this argument can really only be one of four valid statuses that Juju's `status-set` will accept `('active', 'blocked', 'maintenance', 'waiting')`. This new type information is propagated to `StatusBase` and relevant tests, including `_TestingModelBackend`.

Since the arguments to `status_set` are already partially validated,  (`is_app` being a bool in `status_set`, `status` and `message` being strings in `Application.status` and `Unit.status`), and and the valid values for `status` are now encoded and associated with this method, we can check that `status` is valid in `status_set` and raise a `ModelError` immediately if not, rather than catching `status-set`'s exit code and only then raising a `ModelError`.

Tests are updated to account for this (`test_collect_status_priority` split into `test_collect_status_priority_valid` and `test_collect_status_priority_invalid` since `status-set` isn't called in the invalid case, and `test_local_set_invalid_status` is updated to account for this too).

Since `is_app` and `status` are both validated in `status_set`, it makes sense to also validate the type of `message` here, removing redundant checks in `Application.status` and cleaning up a `fixme` in `Unit.status`. A test is added to cover this (`test_status_set_message_not_str_raises`).

Finally, fix a couple of broken unit tests that interact with `StatusBase` and `status_set` (`test_base_status_instance_raises` and `test_status_set_is_app_not_bool_raises` respectively).

dimaqq: ```
  /home/runner/work/operator/operator/src/charm.py:119:33 - error: Argument of type "str" cannot be assigned to parameter "name" of type "_StatusName" in function "from_name"
    Type "str" is not assignable to type "_StatusName"
      "str" is not assignable to type "Literal['active']"
      "str" is not assignable to type "Literal['blocked']"
      "str" is not assignable to type "Literal['maintenance']"
      "str" is not assignable to type "Literal['waiting']"
      "str" is not assignable to type "Literal['error']"
      "str" is not assignable to type "Literal['unknown']" (reportArgumentType)
1 error, 0 warnings, 0 informations 
```

Given that it's an observability job that fails, I suspect that this change is a breaking change subject to how charmers use our code today.

Please check.

P.S. I recommend amazing https://github.com/tonyandrewmeyer/charm-analysis to try new code across some 120~160 charms.
james-garner-canonical: This PR breaks a type checking test in canonical/prometheus-k8s-operator where StatusBase.from_name(name, message) [is called](https://github.com/canonical/prometheus-k8s-operator/blob/27025e5d840ec09331608b398cf4e49e5fe256e9/src/charm.py#L119). The error is because `name` is typed as `str` in this charm, while this PR changes the type of the `name` argument in `from_name` from `str` to `Literal['active', 'blocked', 'maintenance', 'waiting', 'error', 'unknown']`.

Since `StatusBase.from_name` is part of the public api of ops, this isn't good. On the other hand, `from_name` does [currently](https://github.com/canonical/operator/blob/e56686444e3d1504c1785262a1bb3610e1c5182a/ops/model.py#L1888) raises a `KeyError` at runtime if called with a `name` that isn't one of those six.

Solutions here include:

1. change the type of the `name` argument of `from_name` back to `str`, and `cast`/etc internally (it's not our users' problem)
2. require users to either more strongly type their arguments, use a `cast`, or use a `type: ignore[asignment]`
2a. make it easier to use the type `Literal['active', 'blocked', 'maintenance', 'waiting', 'error', 'unknown']` externally (e.g. expose a `model.StatusName` type alias in the public api)
2b. submit a PR to canonical/prometheus-k8s-operator to help with this

I'd lean in favour of (1), and will look into modifying this PR to avoid changes to the type signature of the public api at this point.
james-garner-canonical: > P.S. I recommend amazing https://github.com/tonyandrewmeyer/charm-analysis to try new code across some 120~160 charms.

Thanks, I can see that `StatusBase.from_name` is used in three more charms (see below), as well as in Ben's test-charms repository. I've reverted the signature of this method. I don't think any other changes affect the public API.

Charms that use `StatusBase.from_name`:
```
prometheus-k8s-operator/src/charm.py
sunbeam-charms/ops-sunbeam/ops_sunbeam/compound_status.py
jenkins-k8s-operator/src/charm.py
loki-k8s-operator/src/charm.py
test-charms/statustest/src/statuspool.py
```
dimaqq: Looking at loki example, what they do:

```py
class CompositeStatus(TypedDict):
    """Per-component status holder."""

    # These are going to go into stored state, so we must use marshallable objects.
    # They are passed to StatusBase.from_name().
    k8s_patch: Tuple[str, str]
    config: Tuple[str, str]
    rules: Tuple[str, str]
    retention: Tuple[str, str]

...


    def _on_k8s_patch_failed(self, event: K8sResourcePatchFailedEvent):
        self._stored.status["k8s_patch"] = to_tuple(BlockedStatus(cast(str, event.message)))

...

def to_status(tpl: Tuple[str, str]) -> StatusBase:
    """Convert a tuple to a StatusBase, so it could be used natively with ops."""
    name, message = tpl
    return StatusBase.from_name(name, message)

...

    def _on_collect_unit_status(self, event: CollectStatusEvent):
        # "Pull" statuses
        # TODO refactor _configure to turn the "rules" status into a "pull" status.

        # "Push" statuses
        for status in self._stored.status.values():
            event.add_status(to_status(status))
```

Arguably that code could be refactored, though it's unclear exactly how.

One option would be to change str to a union of literals...

Another option would be observe that `_stored` could hypothetically persist over charm upgrades and therefore ops library upgrades and downgrades, and the run-time check against supported values is therefore needed should the allow list ever change.

In any case, good call to keep this PR small and leave more extensive changes for later / slower process.
james-garner-canonical: Interesting observation about data persisting across library versions. I definitely wouldn't want to remove any runtime checks when adding type checks (though here it would just fail with a relatively traceable KeyError on dictionary access).

The data flow you've pointed out in loki could indeed simply be annotated with `Literal`s, which would make it compatible with both the current api and the breaking change (now reverted). Defining the `Literal['active', 'blocked', 'maintenance', 'waiting', 'error', 'unknown']` is a bit painful for users, and isn't  very 'single source of truth', especially when thinking about changing ops/juju versions, so I wouldn't want to recommend this pattern to library users generally.

I wonder if it would be worth making the `StatusName` (and `SettableStatusName`?) -- type aliases for `Literal`s -- part of the public api, if not in this PR, then perhaps in future.
dimaqq: 115 out of 115 (100%) runs passed. 🚀 

I would have imagined at least some charm to explicitly set 
"error" for good measure, which is now disallowed, but alas! no such thing 😅 
james-garner-canonical: > Generally looks good - a couple of small notes/queries.
> 
> One other thing I noticed while reviewing the docs to see if they fit with this: the `add_status` method accepts both `error` and `unknown`, and the doc lists those. It reads very much like you can add those statuses, but especially for `error` if you include that then it'll definitely fail because it will "win" the evaluation and then can't be set.
> 
> Maybe it makes this PR too big to adjust that, but if so we should probably open a ticket about that too.

I've opened an issue to discuss this (#1356). 
james-garner-canonical: I'm going to defer work on (https://github.com/canonical/operator/issues/1356 -- i.e.  making it an error to call `add_status` with `ErrorStatus` or `UnknownStatus`) to a future PR. Pending finalising the name for the `_StatusName` / `_SettableStatusName` type aliases (see https://github.com/canonical/operator/pull/1354#discussion_r1742903566), I don't anticipate doing further work on this PR in the meantime. There isn't really any rush for getting it finalised, but I would welcome any feedback on the names. For ease of reference, I'll post my two ideas below, but I'm open to suggestions.

```python
_ReadOnlyStatusName = Literal['error', 'unknown']
_SettableStatusName = Literal['active', 'blocked', 'maintenance', 'waiting']
_StatusName = Union[_SettableStatusName, _ReadOnlyStatusName]
_StatusDict = TypedDict('_StatusDict', {'status': _StatusName, 'message': str})
_SETTABLE_STATUS_NAMES: Tuple[_SettableStatusName, ...] = get_args(_SettableStatusName)
```
Prefixes, but slightly more comprehensible in context?

```python
_StatusNameReadOnly = Literal['error', 'unknown']
_StatusNameSettable = Literal['active', 'blocked', 'maintenance', 'waiting']
_StatusName = Union[_StatusNameSettable, _StatusNameReadOnly]
_StatusDict = TypedDict('_StatusDict', {'status': _StatusName, 'message': str})
_SETTABLE_STATUS_NAMES: Tuple[_StatusNameSettable, ...] = get_args(_StatusNameSettable)
```
Are suffixes better here?
james-garner-canonical: I think this is ready for a final review and merge if anyone wants to take another look and provide the second approval (@tonyandrewmeyer, @benhoyt).
james-garner-canonical: Discussed changes have been made, happy to approve and merge @benhoyt ?

Patch:
diff --git a/ops/__init__.py b/ops/__init__.py
index 19679cf..3108647 100644
--- a/ops/__init__.py
+++ b/ops/__init__.py
@@ -162,6 +162,7 @@ __all__ = [  # noqa: RUF022 `__all__` is not sorted
     'SecretRotate',
     'ServiceInfoMapping',
     'StatusBase',
+    'StatusName',
     'Storage',
     'StorageMapping',
     'TooManyRelatedAppsError',
@@ -309,6 +310,7 @@ from .model import (
     SecretRotate,
     ServiceInfoMapping,
     StatusBase,
+    StatusName,
     Storage,
     StorageMapping,
     TooManyRelatedAppsError,
diff --git a/ops/_private/harness.py b/ops/_private/harness.py
index dcabdf3..0933440 100644
--- a/ops/_private/harness.py
+++ b/ops/_private/harness.py
@@ -61,7 +61,7 @@ from ops import charm, framework, model, pebble, storage
 from ops._private import yaml
 from ops.charm import CharmBase, CharmMeta, RelationRole
 from ops.jujucontext import _JujuContext
-from ops.model import Container, RelationNotFoundError, _NetworkDict
+from ops.model import Container, RelationNotFoundError, StatusName, _NetworkDict
 from ops.pebble import ExecProcess
 
 ReadableBuffer = Union[bytes, str, StringIO, BytesIO, BinaryIO]
@@ -80,11 +80,10 @@ _FileKwargs = TypedDict(
 
 _RelationEntities = TypedDict('_RelationEntities', {'app': str, 'units': List[str]})
 
-_StatusName = Literal['unknown', 'blocked', 'active', 'maintenance', 'waiting']
 _RawStatus = TypedDict(
     '_RawStatus',
     {
-        'status': _StatusName,
+        'status': StatusName,
         'message': str,
     },
 )
@@ -2491,7 +2490,7 @@ class _TestingModelBackend:
         else:
             return self._unit_status
 
-    def status_set(self, status: '_StatusName', message: str = '', *, is_app: bool = False):
+    def status_set(self, status: 'StatusName', message: str = '', *, is_app: bool = False):
         if status in [model.ErrorStatus.name, model.UnknownStatus.name]:
             raise model.ModelError(
                 f'ERROR invalid status "{status}", expected one of'
diff --git a/ops/model.py b/ops/model.py
index 7571bd2..335dd9a 100644
--- a/ops/model.py
+++ b/ops/model.py
@@ -42,6 +42,7 @@ from typing import (
     Generator,
     Iterable,
     List,
+    Literal,
     Mapping,
     MutableMapping,
     Optional,
@@ -51,6 +52,7 @@ from typing import (
     Type,
     TypedDict,
     Union,
+    get_args,
 )
 
 import ops
@@ -69,7 +71,11 @@ K8sSpec = Mapping[str, Any]
 _StorageDictType = Dict[str, Optional[List['Storage']]]
 _BindingDictType = Dict[Union[str, 'Relation'], 'Binding']
 
-_StatusDict = TypedDict('_StatusDict', {'status': str, 'message': str})
+_ReadOnlyStatusName = Literal['error', 'unknown']
+_SettableStatusName = Literal['active', 'blocked', 'maintenance', 'waiting']
+StatusName = Union[_SettableStatusName, _ReadOnlyStatusName]
+_StatusDict = TypedDict('_StatusDict', {'status': StatusName, 'message': str})
+_SETTABLE_STATUS_NAMES: Tuple[_SettableStatusName, ...] = get_args(_SettableStatusName)
 
 # mapping from relation name to a list of relation objects
 _RelationMapping_Raw = Dict[str, Optional[List['Relation']]]
@@ -424,9 +430,12 @@ class Application:
         if not self._backend.is_leader():
             raise RuntimeError('cannot set application status as a non-leader unit')
 
-        for _key in {'name', 'message'}:
-            assert isinstance(getattr(value, _key), str), f'status.{_key} must be a string'
-        self._backend.status_set(value.name, value.message, is_app=True)
+        self._backend.status_set(
+            typing.cast(_SettableStatusName, value.name),  # status_set will validate at runtime
+            value.message,
+            is_app=True,
+        )
+
         self._status = value
 
     def planned_units(self) -> int:
@@ -590,8 +599,11 @@ class Unit:
         if not self._is_our_unit:
             raise RuntimeError(f'cannot set status for a remote unit {self}')
 
-        # fixme: if value.messages
-        self._backend.status_set(value.name, value.message, is_app=False)
+        self._backend.status_set(
+            typing.cast(_SettableStatusName, value.name),  # status_set will validate at runtime
+            value.message,
+            is_app=False,
+        )
         self._status = value
 
     def __repr__(self):
@@ -1885,10 +1897,10 @@ class StatusBase:
     directly use the child class such as :class:`ActiveStatus` to indicate their status.
     """
 
-    _statuses: Dict[str, Type['StatusBase']] = {}
+    _statuses: Dict[StatusName, Type['StatusBase']] = {}
 
-    # Subclasses should override this attribute
-    name = ''
+    # Subclasses must provide this attribute
+    name: StatusName
 
     def __init__(self, message: str = ''):
         if self.__class__ is StatusBase:
@@ -1911,7 +1923,8 @@ class StatusBase:
         does not have an associated message.
 
         Args:
-            name: Name of the status, for example "active" or "blocked".
+            name: Name of the status, one of:
+                "active", "blocked", "maintenance", "waiting", "error", or "unknown".
             message: Message to include with the status.
 
         Raises:
@@ -1921,12 +1934,12 @@ class StatusBase:
             # unknown is special
             return UnknownStatus()
         else:
-            return cls._statuses[name](message)
+            return cls._statuses[typing.cast(StatusName, name)](message)
 
     @classmethod
     def register(cls, child: Type['StatusBase']):
         """Register a Status for the child's name."""
-        if not isinstance(child.name, str):
+        if not (hasattr(child, 'name') and isinstance(child.name, str)):
             raise TypeError(
                 f"Can't register StatusBase subclass {child}: ",
                 'missing required `name: str` class attribute',
@@ -1960,7 +1973,7 @@ class UnknownStatus(StatusBase):
     charm has not called status-set yet.
 
     This status is read-only; trying to set unit or application status to
-    ``UnknownStatus`` will raise :class:`ModelError`.
+    ``UnknownStatus`` will raise :class:`InvalidStatusError`.
     """
 
     name = 'unknown'
@@ -1981,7 +1994,7 @@ class ErrorStatus(StatusBase):
     human intervention in order to operate correctly).
 
     This status is read-only; trying to set unit or application status to
-    ``ErrorStatus`` will raise :class:`ModelError`.
+    ``ErrorStatus`` will raise :class:`InvalidStatusError`.
     """
 
     name = 'error'
@@ -3415,13 +3428,15 @@ class _ModelBackend:
         #       status-data: {}
 
         if is_app:
-            content = typing.cast(Dict[str, Dict[str, str]], content)
+            content = typing.cast(Dict[str, '_StatusDict'], content)
             app_status = content['application-status']
             return {'status': app_status['status'], 'message': app_status['message']}
         else:
             return typing.cast('_StatusDict', content)
 
-    def status_set(self, status: str, message: str = '', *, is_app: bool = False) -> None:
+    def status_set(
+        self, status: _SettableStatusName, message: str = '', *, is_app: bool = False
+    ) -> None:
         """Set a status of a unit or an application.
 
         Args:
@@ -3432,6 +3447,10 @@ class _ModelBackend:
         """
         if not isinstance(is_app, bool):
             raise TypeError('is_app parameter must be boolean')
+        if not isinstance(message, str):
+            raise TypeError('message parameter must be a string')
+        if status not in _SETTABLE_STATUS_NAMES:
+            raise InvalidStatusError(f'status must be in {_SETTABLE_STATUS_NAMES}, not {status!r}')
         self._run('status-set', f'--application={is_app}', status, message)
 
     def storage_list(self, name: str) -> List[int]:


Test patch:
diff --git a/test/test_charm.py b/test/test_charm.py
index b6879b1..0939fbb 100644
--- a/test/test_charm.py
+++ b/test/test_charm.py
@@ -22,7 +22,7 @@ import yaml
 
 import ops
 import ops.charm
-from ops.model import ModelError
+from ops.model import ModelError, StatusName
 
 from .test_helpers import FakeScript, create_framework
 
@@ -956,22 +956,20 @@ def test_add_status_type_error(request: pytest.FixtureRequest, fake_script: Fake
 @pytest.mark.parametrize(
     'statuses,expected',
     [
-        (['blocked', 'error'], 'error'),
         (['waiting', 'blocked'], 'blocked'),
         (['waiting', 'maintenance'], 'maintenance'),
         (['active', 'waiting'], 'waiting'),
         (['active', 'unknown'], 'active'),
-        (['unknown'], 'unknown'),
     ],
 )
-def test_collect_status_priority(
+def test_collect_status_priority_valid(
     request: pytest.FixtureRequest,
     fake_script: FakeScript,
-    statuses: typing.List[str],
+    statuses: typing.List[StatusName],
     expected: str,
 ):
     class MyCharm(ops.CharmBase):
-        def __init__(self, framework: ops.Framework, statuses: typing.List[str]):
+        def __init__(self, framework: ops.Framework, statuses: typing.List[StatusName]):
             super().__init__(framework)
             self.framework.observe(self.on.collect_app_status, self._on_collect_status)
             self.statuses = statuses
@@ -991,6 +989,36 @@ def test_collect_status_priority(
     assert status_set_calls == [['status-set', '--application=True', expected, '']]
 
 
+@pytest.mark.parametrize(
+    'statuses',
+    [
+        ['blocked', 'error'],
+        ['unknown'],
+    ],
+)
+def test_collect_status_priority_invalid(
+    request: pytest.FixtureRequest,
+    fake_script: FakeScript,
+    statuses: typing.List[StatusName],
+):
+    class MyCharm(ops.CharmBase):
+        def __init__(self, framework: ops.Framework, statuses: typing.List[StatusName]):
+            super().__init__(framework)
+            self.framework.observe(self.on.collect_app_status, self._on_collect_status)
+            self.statuses = statuses
+
+        def _on_collect_status(self, event: ops.CollectStatusEvent):
+            for status in self.statuses:
+                event.add_status(ops.StatusBase.from_name(status, ''))
+
+    fake_script.write('is-leader', 'echo true')
+
+    framework = create_framework(request)
+    charm = MyCharm(framework, statuses=statuses)
+    with pytest.raises(ModelError):
+        ops.charm._evaluate_status(charm)
+
+
 def test_meta_links():
     # Each type of link can be a single item.
     meta = ops.CharmMeta.from_yaml("""
diff --git a/test/test_model.py b/test/test_model.py
index f8af184..d9e8ac0 100644
--- a/test/test_model.py
+++ b/test/test_model.py
@@ -823,8 +823,8 @@ class TestModel:
         class NoNameStatus(ops.StatusBase):
             pass
 
-        with pytest.raises(AttributeError):
-            ops.StatusBase.register_status(NoNameStatus)  # type: ignore
+        with pytest.raises(TypeError):
+            ops.StatusBase.register(NoNameStatus)
 
     def test_status_repr(self):
         test_cases = {
@@ -2878,34 +2878,28 @@ class TestModelBackend:
                 case()
 
     def test_local_set_invalid_status(self, fake_script: FakeScript):
-        # juju returns exit code 1 if you ask to set status to 'unknown' or 'error'
+        # ops will directly raise InvalidStatusError if you try to set status to unknown or error
         meta = ops.CharmMeta.from_yaml("""
             name: myapp
         """)
         model = ops.Model(meta, self.backend)
-        fake_script.write('status-set', 'exit 1')
         fake_script.write('is-leader', 'echo true')
 
-        with pytest.raises(ops.ModelError):
+        with pytest.raises(ops.InvalidStatusError):
             model.unit.status = ops.UnknownStatus()
-        with pytest.raises(ops.ModelError):
+        with pytest.raises(ops.InvalidStatusError):
             model.unit.status = ops.ErrorStatus()
 
-        assert fake_script.calls(True) == [
-            ['status-set', '--application=False', 'unknown', ''],
-            ['status-set', '--application=False', 'error', ''],
-        ]
+        assert fake_script.calls(True) == []
 
-        with pytest.raises(ops.ModelError):
+        with pytest.raises(ops.InvalidStatusError):
             model.app.status = ops.UnknownStatus()
-        with pytest.raises(ops.ModelError):
+        with pytest.raises(ops.InvalidStatusError):
             model.app.status = ops.ErrorStatus()
 
         # A leadership check is needed for application status.
         assert fake_script.calls(True) == [
             ['is-leader', '--format=json'],
-            ['status-set', '--application=True', 'unknown', ''],
-            ['status-set', '--application=True', 'error', ''],
         ]
 
     @pytest.mark.parametrize('name', ['active', 'waiting', 'blocked', 'maintenance', 'error'])
@@ -2949,9 +2943,20 @@ class TestModelBackend:
         assert model.app.status.message == 'bar'
 
     def test_status_set_is_app_not_bool_raises(self):
-        for is_app_v in [None, 1, 2.0, 'a', b'beef', object]:
+        for is_app_v in [None, 1, 2.0, 'a', b'beef', object()]:
             with pytest.raises(TypeError):
-                self.backend.status_set(ops.ActiveStatus, is_app=is_app_v)  # type: ignore
+                self.backend.status_set(
+                    'active',
+                    is_app=is_app_v,  # type: ignore[assignment]
+                )
+
+    def test_status_set_message_not_str_raises(self):
+        for message in [None, 1, 2.0, True, b'beef', object()]:
+            with pytest.raises(TypeError):
+                self.backend.status_set(
+                    'active',
+                    message=message,  # type: ignore[assignment]
+                )
 
     def test_storage_tool_errors(self, fake_script: FakeScript):
         fake_script.write('storage-list', 'echo fooerror >&2 ; exit 1')