chore: pull _check_idle (single step) from juju#1104

juju/model.py

@@ -61,7 +61,7 @@
 
 if TYPE_CHECKING:
     from .application import Application
-    from .client._definitions import FullStatus
+    from .client._definitions import FullStatus, UnitStatus
     from .constraints import StorageConstraintDict
     from .machine import Machine
     from .relation import Relation
@@ -2655,14 +2655,21 @@ async def get_action_status(self, uuid_or_prefix=None, name=None):
             results[tag.untag("action-", a.action.tag)] = a.status
         return results
 
-    async def get_status(self, filters=None, utc=False) -> FullStatus:
+    async def get_status(self, filters=None, utc: bool = False) -> FullStatus:
         """Return the status of the model.
 
         :param str filters: Optional list of applications, units, or machines
             to include, which can use wildcards ('*').
-        :param bool utc: Display time as UTC in RFC3339 format
+        :param bool utc: Deprecated, display time as UTC in RFC3339 format
 
         """
+        if utc:
+            warnings.warn(
+                "Model.get_status() utc= parameter is deprecated",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+
         client_facade = client.ClientFacade.from_connection(self.connection())
         return await client_facade.FullStatus(patterns=filters)
 
@@ -3204,6 +3211,147 @@ def _raise_for_status(entities: dict[str, list[str]], status: Any):
                 last_log_time = datetime.now()
             await jasyncio.sleep(check_freq)
 
+    async def _check_idle(
+        self,
+        *,
+        apps: list[str],
+        raise_on_error: bool,
+        raise_on_blocked: bool,
+        status: str | None,
+        wait_for_at_least_units: int | None,
+        wait_for_exact_units: int | None,
+        timeout: float | None,
+        idle_period: float,
+        _wait_for_units: int,
+        idle_times: dict[str, datetime],
+        units_ready: set[str],
+        last_log_time: list[datetime | None],
+        start_time: datetime,
+    ) -> bool:
+        now = datetime.now()
+        expected_idle_since = now - timedelta(seconds=idle_period)
+        full_status = await self.get_status()
+        # import pdb; pdb.set_trace()
+
+        # FIXME check this precedence
+        for app_name in apps:
+            if not full_status.applications.get(app_name):
+                logging.info("Waiting for app %r", app_name)
+                return False
+
+        # Order of errors:
+        #
+        # Machine error (any unit of any app from apps)
+        # Agent error (-"-)
+        # Workload error (-"-)
+        # App error (any app from apps)
+        #
+        # Workload blocked (any unit of any app from apps)
+        # App blocked (any app from apps)
+        units: dict[str, UnitStatus] = {}
+
+        for app_name in apps:
+            # assert full_status.applications[app_name]
+            app = full_status.applications[app_name]
+            assert app
+            for unit_name, unit in app.units.items():
+                assert unit
+                units[unit_name] = unit
+
+        for unit_name, unit in units.items():
+            if unit.machine:
+                machine = full_status.machines[unit.machine]
+                assert machine
+                assert machine.instance_status
+                if machine.instance_status.status == "error" and raise_on_error:
+                    raise JujuMachineError(
+                        f"{unit_name!r} machine {unit.machine!r} has errored: {machine.instance_status.info!r}"
+                    )
+
+        for unit_name, unit in units.items():
+            assert unit.agent_status
+            if unit.agent_status.status == "error" and raise_on_error:
+                raise JujuAgentError(
+                    f"{unit_name!r} agent has errored: {unit.agent_status.info!r}"
+                )
+
+        for unit_name, unit in units.items():
+            assert unit.workload_status
+            if unit.workload_status.status == "error" and raise_on_error:
+                raise JujuUnitError(
+                    f"{unit_name!r} workload has errored: {unit.workload_status.info!r}"
+                )
+
+        for app_name in apps:
+            app = full_status.applications[app_name]
+            assert app
+            assert app.status
+            if app.status.status == "error" and raise_on_error:
+                raise JujuAppError(f"{app_name!r} has errored: {app.status.info!r}")
+
+        for unit_name, unit in units.items():
+            assert unit.workload_status
+            if unit.workload_status.status == "blocked" and raise_on_blocked:
+                raise JujuUnitError(
+                    f"{unit_name!r} workload is blocked: {unit.workload_status.info!r}"
+                )
+
+        for app_name in apps:
+            app = full_status.applications[app_name]
+            assert app
+            assert app.status
+            if app.status.status == "blocked" and raise_on_blocked:
+                raise JujuAppError(f"{app_name!r} is blocked: {app.status.info!r}")
+
+        for unit_name, unit in units.items():
+            assert unit.agent_status
+            idle_times.setdefault(unit_name, now)
+            if unit.agent_status.status != "idle":
+                idle_times[unit_name] = now
+
+        for app_name in apps:
+            ready_units = []
+            app = full_status.applications[app_name]
+            assert app
+            for unit in app.units.values():
+                assert unit
+                assert unit.agent_status
+                assert unit.workload_status
+
+                if unit.agent_status.status != "idle":
+                    continue
+                if status and unit.workload_status.status != status:
+                    continue
+
+                ready_units.append(unit)
+
+            if wait_for_exact_units is None and len(ready_units) < _wait_for_units:
+                logging.info(
+                    "Waiting for app %r units %s/%s",
+                    app_name,
+                    len(ready_units),
+                    _wait_for_units,
+                )
+                return False
+
+            if (
+                wait_for_exact_units is not None
+                and len(ready_units) != wait_for_exact_units
+            ):
+                logging.info(
+                    "Waiting for app %r units %s/%s",
+                    app_name,
+                    len(ready_units),
+                    _wait_for_units,
+                )
+                return False
+
+        if busy := [n for n, t in idle_times.items() if expected_idle_since < t]:
+            logging.info("Waiting for %s to be idle enough", busy)
+            return False
+
+        return True
+
 
 def _create_consume_args(offer, macaroon, controller_info):
     """Convert a typed object that has been normalised to a overridden typed