Repository: canonical/operator
Pull request #1379 (2024-09-18 00:31:38)

refactor: keep the unittest.mock names in the 'mock' namespace
No functional changes (including to the tests) - this just adjusts the imports so that the names from `unittest.mock` are imported into a `mock` namespace, rather than all into the module namespace.

Primarily, this ensures that `typing.Any` and `unittest.mock.ANY` are clearly distinguished by more than case, but it's also generally tidier ("Namespaces are one honking great idea").

See discussion in #1373.


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..0eadd4d 100644
--- a/test/test_model.py
+++ b/test/test_model.py
@@ -24,7 +24,7 @@ import typing
 import unittest
 from collections import OrderedDict
 from textwrap import dedent
-from unittest.mock import ANY, MagicMock, patch
+from unittest import mock
 
 import pytest
 
@@ -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 = {
@@ -1114,9 +1114,9 @@ class TestModel:
         assert str(excinfo.value) == 'ERROR cannot get status\n'
         assert excinfo.value.args[0] == 'ERROR cannot get status\n'
 
-    @patch('grp.getgrgid')
-    @patch('pwd.getpwuid')
-    def test_push_path_unnamed(self, getpwuid: MagicMock, getgrgid: MagicMock):
+    @mock.patch('grp.getgrgid')
+    @mock.patch('pwd.getpwuid')
+    def test_push_path_unnamed(self, getpwuid: mock.MagicMock, getgrgid: mock.MagicMock):
         getpwuid.side_effect = KeyError
         getgrgid.side_effect = KeyError
         harness = ops.testing.Harness(
@@ -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')
@@ -3364,7 +3369,9 @@ class TestSecrets:
         assert secret.id == 'secret:123'
         assert secret.label is None
 
-        assert fake_script.calls(clear=True) == [['secret-add', '--owner', 'application', ANY]]
+        assert fake_script.calls(clear=True) == [
+            ['secret-add', '--owner', 'application', mock.ANY]
+        ]
         assert fake_script.secrets() == {'foo': 'x'}
 
     def test_app_add_secret_args(self, fake_script: FakeScript, model: ops.Model):
@@ -3395,8 +3402,8 @@ class TestSecrets:
                 'hourly',
                 '--owner',
                 'application',
-                ANY,
-                ANY,
+                mock.ANY,
+                mock.ANY,
             ]
         ]
         assert fake_script.secrets() == {'foo': 'x', 'bar': 'y'}
@@ -3409,7 +3416,7 @@ class TestSecrets:
         assert secret.id == 'secret:345'
         assert secret.label is None
 
-        assert fake_script.calls(clear=True) == [['secret-add', '--owner', 'unit', ANY]]
+        assert fake_script.calls(clear=True) == [['secret-add', '--owner', 'unit', mock.ANY]]
         assert fake_script.secrets() == {'foo': 'x'}
 
     def test_unit_add_secret_args(self, fake_script: FakeScript, model: ops.Model):
@@ -3440,8 +3447,8 @@ class TestSecrets:
                 'yearly',
                 '--owner',
                 'unit',
-                ANY,
-                ANY,
+                mock.ANY,
+                mock.ANY,
             ]
         ]
         assert fake_script.secrets() == {'foo': 'w', 'bar': 'z'}
@@ -3771,9 +3778,9 @@ class TestSecretClass:
             secret.set_content({'s': 't'})  # ensure it validates content (key too short)
 
         assert fake_script.calls(clear=True) == [
-            ['secret-set', f'secret://{model._backend.model_uuid}/x', ANY],
+            ['secret-set', f'secret://{model._backend.model_uuid}/x', mock.ANY],
             ['secret-info-get', '--label', 'y', '--format=json'],
-            ['secret-set', f'secret://{model._backend.model_uuid}/z', ANY],
+            ['secret-set', f'secret://{model._backend.model_uuid}/z', mock.ANY],
         ]
         assert fake_script.secrets() == {'foo': 'bar', 'bar': 'foo'}