chore: more app and unit counting mechanics

juju/model.py

@@ -63,7 +63,7 @@
 from ._sync import SyncCacheLine, ThreadedAsyncRunner
 
 if TYPE_CHECKING:
-    from .client._definitions import FullStatus
+    from .client._definitions import ApplicationStatus, FullStatus
     from .application import Application
     from .machine import Machine
     from .relation import Relation
@@ -3047,15 +3047,15 @@ async def _check_idle(
         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(),
+        idle_times: Dict[str, datetime],
+        units_ready: Set[str],
+        last_log_time: List[Optional[datetime]],
+        start_time: datetime,
     ) -> bool:
         now = datetime.now()
         expected_idle_since = now - timedelta(seconds=idle_period)
         full_status = await self.get_status()
-        __import__("pdb").set_trace()
+        # __import__("pdb").set_trace()
 
         for app_name in apps:
             if not (app := full_status.applications.get(app_name)):
@@ -3083,13 +3083,73 @@ async def _check_idle(
                              app_name, len(app.units), _wait_for_units)
                 return False
 
+
             # TODO: refactor to simplify later
             # datetime -> float; check vs outer loop
             idle_since = idle_times.setdefault(app_name, now)
             if expected_idle_since < idle_since:
                 logging.info("App %r has not been idle long enough", app_name)
                 return False
-            # TODO continue here...
+
+            ready_units = [
+                self._check_idle_unit(
+                    unit_name,
+                    app_name=app_name,
+                    full_status=full_status,
+                    raise_on_error=raise_on_error,
+                    raise_on_blocked=raise_on_blocked,
+                    status=status,
+                    idle_period=idle_period,
+                    idle_times=idle_times,
+                    )
+                for unit_name in app.units
+            ]
+
+        return True
+
+    async def _check_idle_unit(
+        self,
+        unit_name: str,
+        *,
+        app_name: str,
+        full_status: FullStatus,
+        raise_on_error: bool,
+        raise_on_blocked: bool,
+        status: Optional[str],
+        #
+        idle_period: float,
+        idle_times: Dict[str, datetime],
+    ):
+        app = full_status.applications[app_name]
+        assert app
+        unit = app.units[unit_name]
+        assert unit
+        machines = full_status.machines
+        assert machines is not None
+
+        assert unit.agent_status
+        assert unit.workload_status
+        statuses = (unit.agent_status, unit.workload_status)
+        status_text = f"agent {unit.agent_status.status!r} {unit.agent_status.info!r} workload {unit.workload_status.status!r} {unit.workload_status.info!r}"
+
+        if raise_on_error and "error" in [s.status for s in statuses]:
+            raise JujuUnitError(f"Unit {unit_name!r} is in error: {status_text}")
+
+        if raise_on_error and "blocked" in [s.status for s in statuses]:
+            raise JujuUnitError(f"Unit {unit_name!r} is blocked: {status_text}")
+
+        if unit.machine:
+            machine = machines[unit.machine]
+            assert machine
+            # reveal_type(machine)
+            # FIXME: which of these must be covered?
+            # reveal_type(machine.agent_status)
+            # reveal_type(machine.instance_status)
+            # reveal_type(machine.modification_status)
+            raise NotImplementedError("Need FullStatus sample for machines")
+
+        # TODO implement timers
+
         return True
 
     async def _legacy_check_idle(
@@ -3104,12 +3164,12 @@ async def _legacy_check_idle(
         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(),
+        idle_times: Dict[str, datetime],
+        units_ready: Set[str],
+        last_log_time: List[Optional[datetime]],  # = [None]
+        start_time: datetime,
     ):
-        __import__("pdb").set_trace()
+        # __import__("pdb").set_trace()
         _timeout = timedelta(seconds=timeout) if timeout is not None else None
         _idle_period = timedelta(seconds=idle_period)
         log_interval = timedelta(seconds=30)

tests/unit/test_wait_for_idle.py

@@ -118,6 +118,10 @@ def kwargs() -> Dict[str, Any]:
         timeout=100,
         idle_period=0,
         _wait_for_units=1,
+        idle_times={},
+        units_ready=set(),
+        last_log_time=[None],
+        start_time=datetime.now(),
     )