wip

juju/application.py

@@ -27,6 +27,7 @@
 from .status import derive_status
 from ._sync import cache_until_await
 from .url import URL
+from .unit import Unit
 from .utils import block_until
 from .version import DEFAULT_ARCHITECTURE
 
@@ -141,7 +142,7 @@ def on_unit_remove(self, callable_):
             callable_, 'unit', 'remove', self._unit_match_pattern)
 
     @property
-    def units(self):
+    def units(self) -> List[Unit]:
         # FIXME need a live call to query units of a given app
         return [
             unit for unit in self.model.units.values()

juju/model.py

@@ -25,11 +25,13 @@
     Dict,
     List,
     Literal,
+    Mapping,
     Optional,
+    overload,
     Set,
     TypeVar,
     TYPE_CHECKING,
-    Union,
+    # Union,
 )
 
 import yaml
@@ -61,6 +63,10 @@
 
 if TYPE_CHECKING:
     from .application import Application
+    from .machine import Machine
+    from .relation import Relation
+    from .remoteapplication import ApplicationOffer, RemoteApplication
+    from .unit import Unit
 
 log = logging.getLogger(__name__)
 
@@ -141,7 +147,27 @@ def __init__(self, model):
         self.model = model
         self.state = dict()
 
-    def _live_entity_map(self, entity_type) -> dict[str, ModelEntity]:
+    @overload
+    def _live_entity_map(self, entity_type: Literal["application"]) -> Dict[str, Application]: ...
+
+    @overload
+    def _live_entity_map(self, entity_type: Literal["applicationOffer"]) -> Dict[str, ApplicationOffer]: ...
+
+    @overload
+    def _live_entity_map(self, entity_type: Literal["machine"]) -> Dict[str, Machine]: ...
+
+    @overload
+    def _live_entity_map(self, entity_type: Literal["relation"]) -> Dict[str, Relation]: ...
+
+    @overload
+    def _live_entity_map(self, entity_type: Literal["remoteApplication"]) -> Dict[str, RemoteApplication]: ...
+
+    @overload
+    def _live_entity_map(self, entity_type: Literal["unit"]) -> Dict[str, Unit]: ...
+
+    # FIXME and all the other types
+
+    def _live_entity_map(self, entity_type: str) -> Mapping[str, ModelEntity]:
         """Return an id:Entity map of all the living entities of
         type ``entity_type``.
 
@@ -161,30 +187,30 @@ def applications(self) -> dict[str, Application]:
         return self._live_entity_map('application')
 
     @property
-    def remote_applications(self):
+    def remote_applications(self) -> Dict[str, RemoteApplication]:
         """Return a map of application-name:Application for all remote
         applications currently in the model.
 
         """
         return self._live_entity_map('remoteApplication')
 
     @property
-    def application_offers(self):
+    def application_offers(self) -> Dict[str, ApplicationOffer]:
         """Return a map of application-name:Application for all applications
         offers currently in the model.
         """
         return self._live_entity_map('applicationOffer')
 
     @property
-    def machines(self):
+    def machines(self) -> Dict[str, Machine]:  # FIXME validate that key is in fact a string
         """Return a map of machine-id:Machine for all machines currently in
         the model.
 
         """
         return self._live_entity_map('machine')
 
     @property
-    def units(self):
+    def units(self) -> Dict[str, Unit]:
         """Return a map of unit-id:Unit for all units currently in
         the model.
 
@@ -197,7 +223,7 @@ def subordinate_units(self):
         return {u_name: u for u_name, u in self.units.items() if u.is_subordinate}
 
     @property
-    def relations(self):
+    def relations(self) -> Dict[str, Relation]:
         """Return a map of relation-id:Relation for all relations currently in
         the model.
 
@@ -241,10 +267,12 @@ def apply_delta(self, delta):
             history.append(None)
 
         entity = self.get_entity(delta.entity, delta.get_id())
+        assert entity
         return entity.previous(), entity
 
+    # FIXME this function may explicitly return None, but is the rest of the code prepared for that?
     def get_entity(
-            self, entity_type, entity_id, history_index=-1, connected=True) -> ModelEntity:
+            self, entity_type, entity_id, history_index=-1, connected=True) -> Optional[ModelEntity]:
         """Return an object instance for the given entity_type and id.
 
         By default the object state matches the most recent state from
@@ -1121,30 +1149,30 @@ def applications(self) -> dict[str, Application]:
         return self.state.applications
 
     @property
-    def remote_applications(self):
+    def remote_applications(self) -> Dict[str, RemoteApplication]:
         """Return a map of application-name:Application for all remote
         applications currently in the model.
 
         """
         return self.state.remote_applications
 
     @property
-    def application_offers(self):
+    def application_offers(self) -> Dict[str, ApplicationOffer]:
         """Return a map of application-name:Application for all applications
         offers currently in the model.
         """
         return self.state.application_offers
 
     @property
-    def machines(self):
+    def machines(self) -> Dict[str, Machine]:  # FIXME validate that key is string and not an int
         """Return a map of machine-id:Machine for all machines currently in
         the model.
 
         """
         return self.state.machines
 
     @property
-    def units(self):
+    def units(self) -> Dict[str, Unit]:
         """Return a map of unit-id:Unit for all units currently in
         the model.
 
@@ -2948,21 +2976,59 @@ async def wait_for_idle(self,
                 'Invalid value for wait_for_exact_units : %s' % wait_for_exact_units
 
         while True:
-            if (await self._check_idle(
-                apps=apps,
-                raise_on_error=raise_on_error,
-                raise_on_blocked=raise_on_blocked,
-                status=status,
-                wait_for_at_least_units=wait_for_at_least_units,
-                wait_for_exact_units=wait_for_exact_units,
-                timeout=timeout,
-                idle_period=idle_period,
-                _wait_for_units=_wait_for_units,
-                idle_times=idle_times,
-                units_ready=units_ready,
-                last_log_time=last_log_time,
-                start_time=start_time,
-            )):
+            exc: Optional[Exception] = None
+            legacy_exc: Optional[Exception] = None
+            idle = legacy_idle = False
+            try:
+                idle = await self._check_idle(
+                    apps=apps,
+                    raise_on_error=raise_on_error,
+                    raise_on_blocked=raise_on_blocked,
+                    status=status,
+                    wait_for_at_least_units=wait_for_at_least_units,
+                    wait_for_exact_units=wait_for_exact_units,
+                    timeout=timeout,
+                    idle_period=idle_period,
+                    _wait_for_units=_wait_for_units,
+                    idle_times=idle_times,
+                    units_ready=units_ready,
+                    last_log_time=last_log_time,
+                    start_time=start_time,
+                )
+            except Exception as e:
+                exc = e
+
+            try:
+                legacy_idle = await self._legacy_check_idle(
+                    apps=apps,
+                    raise_on_error=raise_on_error,
+                    raise_on_blocked=raise_on_blocked,
+                    status=status,
+                    wait_for_at_least_units=wait_for_at_least_units,
+                    wait_for_exact_units=wait_for_exact_units,
+                    timeout=timeout,
+                    idle_period=idle_period,
+                    _wait_for_units=_wait_for_units,
+                    idle_times=idle_times,
+                    units_ready=units_ready,
+                    last_log_time=last_log_time,
+                    start_time=start_time,
+                )
+            except Exception as e:
+                legacy_exc = e
+
+            if bool(exc) ^ bool(legacy_exc):
+                warnings.warn(f"Idle loop mismatch: {[exc, legacy_exc]}")
+
+            if exc:
+                raise exc
+            if legacy_exc:
+                raise legacy_exc
+
+            if idle ^ legacy_idle:
+                warnings.warn(f"Idle loop mismatch: {[idle, legacy_idle]}")
+
+            if idle or legacy_idle:
                 return
 
             await jasyncio.sleep(check_freq)
@@ -2983,6 +3049,25 @@ async def _check_idle(
         units_ready: Set[str] = set(),  # The units that are in the desired state
         last_log_time: List[Optional[datetime]] = [None],
         start_time: datetime = datetime.now(),
+    ):
+        return True
+
+    async def _legacy_check_idle(
+        self,
+        *,
+        apps: List[str],
+        raise_on_error: bool,
+        raise_on_blocked: bool,
+        status: Optional[str],
+        wait_for_at_least_units: Optional[int],
+        wait_for_exact_units: Optional[int],
+        timeout: Optional[float],
+        idle_period: float,
+        _wait_for_units: int,
+        idle_times: Dict[str, datetime] = {},
+        units_ready: Set[str] = set(),  # The units that are in the desired state
+        last_log_time: List[Optional[datetime]] = [None],
+        start_time: datetime = datetime.now(),
     ):
         _timeout = timedelta(seconds=timeout) if timeout is not None else None
         _idle_period = timedelta(seconds=idle_period)

tests/unit/test_wait_for_idle.py

@@ -1,17 +1,88 @@
+# Copyright 2024 Canonical Ltd.
+# Licensed under the Apache V2, see LICENCE file for details.
+from __future__ import annotations
+
 import json
 import pytest
+from typing import Dict, List
 from typing import reveal_type as reveal_type
 
+from juju.application import Application
 from juju.client.facade import _convert_response
 from juju.client._definitions import FullStatus
+from juju.model import Model
+from juju.unit import Unit
 
 
 @pytest.fixture
 def full_status_response(pytestconfig: pytest.Config) -> dict:
     return json.loads(((pytestconfig.rootpath / "fullstatus.json")).read_text())
 
 
-def test_foo(full_status_response):
+def test_foo(full_status_response: dict):
     # tweak the response here
     fs = _convert_response(full_status_response, cls=FullStatus)
-    reveal_type(fs)
+
+    assert fs.applications
+    assert fs.applications["hexanator"]
+
+    hexanator = fs.applications["hexanator"]
+    assert hexanator.status  # DetailedStatus
+    assert hexanator.status.status == "active"
+    assert hexanator.status.info == ""
+
+    assert fs.applications["grafana-agent-k8s"]
+    grafana = fs.applications["grafana-agent-k8s"]  # ApplicationStatus
+
+    assert grafana.units["grafana-agent-k8s/0"]
+    grafana0 = grafana.units["grafana-agent-k8s/0"]
+
+    assert grafana0.workload_status  # DetailedStatus
+    assert isinstance(grafana0.workload_status.info, str)
+
+    assert grafana0.workload_status.info.startswith("Missing")
+
+
+class ModelFake(Model):
+    _applications: Dict[str, Application]
+
+    @property
+    def applications(self) -> Dict[str, Application]:
+        return self._applications
+
+    def __init__(self):
+        super().__init__()
+        self._applications = {}
+
+
+class ApplicationFake(Application):
+    _status: str = ""
+    _status_message: str = ""
+    _units: List[Unit]
+
+    @property
+    def status(self) -> str:
+        return self._status
+
+    @property
+    def status_message(self) -> str:
+        return self._status_message
+
+    @property
+    def units(self) -> List[Unit]:
+        return self._units
+
+async def test_something(full_status_response):
+    m = Model()
+    rv = await m._legacy_check_idle(
+            apps=[],
+            raise_on_error=False,
+            raise_on_blocked=False,
+            status=None,
+            wait_for_at_least_units=None,
+            wait_for_exact_units=None,
+            timeout=100,
+            idle_period=100,
+            _wait_for_units=1,
+            )
+    assert rv == 42