chore: refactor unit tests

juju/model/__init__.py

@@ -15,14 +15,15 @@
 import stat
 import sys
 import tempfile
+import time
 import warnings
 import weakref
 import zipfile
 from concurrent.futures import CancelledError
 from datetime import datetime, timedelta
 from functools import partial
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Literal, Mapping, Sequence, overload
+from typing import TYPE_CHECKING, Any, Iterable, Literal, Mapping, overload
 
 import websockets
 import yaml
@@ -3020,7 +3021,7 @@ async def _get_source_api(self, url):
 
     async def wait_for_idle(
         self,
-        apps: Sequence[str] | None = None,
+        apps: Iterable[str] | None = None,
         raise_on_error: bool = True,
         raise_on_blocked: bool = False,
         wait_for_active: bool = False,
@@ -3033,7 +3034,7 @@ async def wait_for_idle(
     ) -> None:
         """Wait for applications in the model to settle into an idle state.
 
-        :param List[str] apps: Optional list of specific app names to wait on.
+        :param Iterable[str]|None apps: Optional list of specific app names to wait on.
             If given, all apps must be present in the model and idle, while other
             apps in the model can still be busy. If not given, all apps currently
             in the model must be idle.
@@ -3108,15 +3109,12 @@ async def wait_for_idle(
         start_time = datetime.now()
 
         if isinstance(apps, (str, bytes, bytearray, memoryview)):
-            raise ValueError("apps must be a Sequence[str] and not a plain type")
+            raise TypeError(f"apps must be a Iterable[str], got {apps=}")
 
-        apps_ = apps or list(self.applications)
-
-        if not isinstance(apps_, collections.abc.Sequence):
-            raise ValueError("apps must be a Sequence[str]")
+        apps_ = list(apps or self.applications)
 
         if any(not isinstance(o, str) for o in apps_):
-            raise ValueError("apps must be a Sequence[str]")
+            raise TypeError(f"apps must be a Iterable[str], got {apps_=}")
 
         idle_times: dict[str, datetime] = {}
         units_ready: set[str] = set()  # The units that are in the desired state
@@ -3248,17 +3246,18 @@ def _raise_for_status(entities: dict[str, list[str]], status: Any):
             _raise_for_status(blocks, "blocked")
             if not busy:
                 break
-            busy = "\n  ".join(busy)
-            if timeout is not None and datetime.now() - start_time > timeout:
-                raise asyncio.TimeoutError("Timed out waiting for model:\n" + busy)
+            if timeout_ is not None and datetime.now() - start_time > timeout_:
+                raise asyncio.TimeoutError(
+                    "\n  ".join(["Timed out waiting for model:", *busy])
+                )
             if last_log_time is None or datetime.now() - last_log_time > log_interval:
-                log.info("Waiting for model:\n  " + busy_)
+                log.info("\n  ".join(["Waiting for model:", *busy]))
                 last_log_time = datetime.now()
             await asyncio.sleep(check_freq)
 
     async def new_wait_for_idle(
         self,
-        apps: Sequence[str] | None = None,
+        apps: Iterable[str] | None = None,
         raise_on_error: bool = True,
         raise_on_blocked: bool = False,
         wait_for_active: bool = False,
@@ -3292,41 +3291,37 @@ async def new_wait_for_idle(
         )
 
         if isinstance(apps, (str, bytes, bytearray, memoryview)):
-            raise ValueError("apps must be a Sequence[str] and not a plain type")
-
-        apps = apps or list(self.applications)
+            raise TypeError(f"apps must be a Iterable[str], got {apps=}")
 
-        if not isinstance(apps, collections.abc.Sequence):
-            raise ValueError("apps must be a Sequence[str]")
+        apps = frozenset(apps or self.applications)
 
         if any(not isinstance(o, str) for o in apps):
-            raise ValueError("apps must be a Sequence[str]")
+            raise TypeError(f"apps must be a Iterable[str], got {apps=}")
 
-        idle_loop = idle._loop(
-            apps=apps,
-            timeout=timeout,
-            wait_for_exact_units=wait_for_exact_units,
-            wait_for_units=wait_for_units,
-            idle_period=idle_period,
-        )
-        next(idle_loop)  # Prime the generator
+        deadline = None if timeout is None else time.monotonic() + timeout
 
-        while True:
-            full_status = await self.get_status()
-
-            ustat = idle._check(
-                full_status,
+        async def status_on_demand():
+            yield idle._check(
+                await self.get_status(),
                 apps=apps,
                 raise_on_error=raise_on_error,
                 raise_on_blocked=raise_on_blocked,
                 status=status,
             )
 
-            yrv = idle_loop.send(ustat)
-
-            if yrv:
+        async for done in idle._loop(
+            status_on_demand(),
+            apps=apps,
+            wait_for_exact_units=wait_for_exact_units,
+            wait_for_units=wait_for_units,
+            idle_period=idle_period,
+        ):
+            if done:
                 break
 
+            if deadline and time.monotonic() > deadline:
+                raise asyncio.TimeoutError(f"Timed out after {timeout}s")
+
             await asyncio.sleep(check_freq)
 
 

juju/model/idle.py

@@ -4,11 +4,10 @@
 
 from __future__ import annotations
 
-import asyncio
 import logging
 import time
 from dataclasses import dataclass
-from typing import Generator, Sequence
+from typing import AsyncIterable
 
 from ..client._definitions import (
     ApplicationStatus,
@@ -26,59 +25,51 @@ class _CheckStatus:
     """Return type for a single interaction, that is _check()."""
 
     units: set[str]
+    """All units visible at this point."""
     ready_units: set[str]
+    """Units in good status (workload, agent, machine?)."""
+    idle_units: set[str]
+    """Units with stable agent status (FIXME details)."""
 
 
-def _loop(
+async def _loop(
+    foo: AsyncIterable[_CheckStatus | None],
     *,
-    apps: Sequence[str],
-    timeout: float | None,
+    apps: frozenset[str],
     wait_for_exact_units: int | None = None,
     wait_for_units: int,
     idle_period: float,
-) -> Generator[bool, _CheckStatus | None, None]:
+) -> AsyncIterable[bool]:
     """The outer, time-dependents logic of a wait_for_idle loop."""
-    status: _CheckStatus | None = yield False  # None at first
-    deadline = None if timeout is None else time.monotonic() + timeout
     idle_since: dict[str, float] = {}
 
-    while True:
+    async for status in foo:
         logger.warning("FIXME unit test debug %r", status)
         now = time.monotonic()
-        if deadline and now > deadline:
-            raise asyncio.TimeoutError(f"Timed out after {now - deadline}")
 
         if not status:
-            status = yield False
+            yield False
             continue
 
         expected_idle_since = now - idle_period
         rv = True
 
-        for name in status.units:
-            idle_since.setdefault(name, now)
-
         # FIXME there's some confusion about what a "busy" unit is
-        # I've simplified the logic to: existing unit is either busy or ready
-        #
-        # However, the "idle_since" check may need to be updated
-        # Right now, a "busy" unit is stamped with ready := since right now
-        #
-        # This is fine for idle_period > 0
-        # But may be problematic for idle_period===0 and wait_for_units==0
-        # (there's a single use case for that in the wild)
-        #
-        # Specifically, a unit in an error state is not ready
-        # But it gets stamped with "now" and seems "idle enough"
-        for name in status.units - status.ready_units:
-            idle_since[name] = now
+        # are we ready when over the last idle_period, every time sampled:
+        # a. >=N units were ready (possibly different each time), or
+        # b. >=N units were ready each time
+        for name in status.units:
+            if name in status.ready_units:
+                idle_since[name] = min(now, idle_since.get(name, float("inf")))
+            else:
+                idle_since[name] = float("inf")
 
         for app_name in apps:
             ready_units = [
                 n for n in status.ready_units if n.startswith(f"{app_name}/")
             ]
             if len(ready_units) < wait_for_units:
-                logger.info(
+                logger.warn(
                     "Waiting for app %r units %s >= %s",
                     app_name,
                     len(status.ready_units),
@@ -90,7 +81,7 @@ def _loop(
                 wait_for_exact_units is not None
                 and len(ready_units) != wait_for_exact_units
             ):
-                logger.info(
+                logger.warn(
                     "Waiting for app %r units %s == %s",
                     app_name,
                     len(ready_units),
@@ -102,16 +93,16 @@ def _loop(
         # Assume that we've got some units ready and some busy
         # What are the semantics for returning True?
         if busy := [n for n, t in idle_since.items() if t > expected_idle_since]:
-            logger.info("Waiting for %s to be idle enough", busy)
+            logger.warn("Waiting for %s to be idle enough", busy)
             rv = False
 
-        status = yield rv
+        yield rv
 
 
 def _check(
     full_status: FullStatus,
     *,
-    apps: Sequence[str],
+    apps: frozenset[str],
     raise_on_error: bool,
     raise_on_blocked: bool,
     status: str | None,
@@ -181,7 +172,7 @@ def _check(
         if app.status.status == "blocked" and raise_on_blocked:
             raise JujuAppError(f"{app_name!r} is blocked: {app.status.info!r}")
 
-    rv = _CheckStatus(set(), set())
+    rv = _CheckStatus(set(), set(), set())
 
     for app_name in apps:
         ready_units = []
@@ -200,6 +191,8 @@ def _check(
             ready_units.append(unit)
             rv.ready_units.add(unit_name)
 
+    # FIXME
+    # rv.idle_units -- depends on agent status only, not workload status
     return rv
 
 

tests/unit/test_idle_check.py

@@ -32,6 +32,7 @@ def test_check_status(full_status: FullStatus, kwargs):
             "mysql-test-app/0",
             "mysql-test-app/1",
         },
+        set(),
     )
 
 
@@ -44,13 +45,13 @@ def test_check_status_missing_app(full_status: FullStatus, kwargs):
 def test_check_status_is_selective(full_status: FullStatus, kwargs):
     kwargs["apps"] = ["hexanator"]
     status = _check(full_status, **kwargs)
-    assert status == _CheckStatus({"hexanator/0"}, {"hexanator/0"})
+    assert status == _CheckStatus({"hexanator/0"}, {"hexanator/0"}, set())
 
 
 def test_no_apps(full_status: FullStatus, kwargs):
     kwargs["apps"] = []
     status = _check(full_status, **kwargs)
-    assert status == _CheckStatus(set(), set())
+    assert status == _CheckStatus(set(), set(), set())
 
 
 def test_missing_app(full_status: FullStatus, kwargs):
@@ -63,7 +64,7 @@ def test_no_units(response: dict[str, Any], kwargs):
     response["response"]["applications"]["hexanator"]["units"].clear()
     kwargs["apps"] = ["hexanator"]
     status = _check(convert(response), **kwargs)
-    assert status == _CheckStatus(set(), set())
+    assert status == _CheckStatus(set(), set(), set())
 
 
 def test_app_error(response: dict[str, Any], kwargs):
@@ -88,7 +89,7 @@ def test_exact_count(response: dict[str, Any], kwargs):
 
     status = _check(convert(response), **kwargs)
     assert status == _CheckStatus(
-        {"hexanator/0", "hexanator/1"}, {"hexanator/0", "hexanator/1"}
+        {"hexanator/0", "hexanator/1"}, {"hexanator/0", "hexanator/1"}, set()
     )
 
 
@@ -98,14 +99,17 @@ def test_ready_units(full_status: FullStatus, kwargs):
     assert status == _CheckStatus(
         {"mysql-test-app/0", "mysql-test-app/1"},
         {"mysql-test-app/0", "mysql-test-app/1"},
+        set(),
     )
 
 
 def test_active_units(full_status: FullStatus, kwargs):
     kwargs["apps"] = ["mysql-test-app"]
     kwargs["status"] = "active"
     status = _check(full_status, **kwargs)
-    assert status == _CheckStatus({"mysql-test-app/0", "mysql-test-app/1"}, set())
+    assert status == _CheckStatus(
+        {"mysql-test-app/0", "mysql-test-app/1"}, set(), set()
+    )
 
 
 def test_ready_unit_requires_idle_agent(response: dict[str, Any], kwargs):
@@ -131,7 +135,9 @@ def test_ready_unit_requires_workload_status(response: dict[str, Any], kwargs):
     kwargs["status"] = "active"
 
     status = _check(convert(response), **kwargs)
-    assert status == _CheckStatus({"hexanator/0", "hexanator/1"}, {"hexanator/0"})
+    assert status == _CheckStatus(
+        {"hexanator/0", "hexanator/1"}, {"hexanator/0"}, set()
+    )
 
 
 def test_agent_error(response: dict[str, Any], kwargs):
@@ -181,7 +187,7 @@ def test_machine_ok(response: dict[str, Any], kwargs):
     kwargs["raise_on_error"] = True
 
     status = _check(convert(response), **kwargs)
-    assert status == _CheckStatus({"hexanator/0"}, {"hexanator/0"})
+    assert status == _CheckStatus({"hexanator/0"}, {"hexanator/0"}, set())
 
 
 def test_machine_error(response: dict[str, Any], kwargs):
@@ -243,7 +249,7 @@ def test_maintenance(response: dict[str, Any], kwargs):
     kwargs["status"] = "maintenance"
 
     status = _check(convert(response), **kwargs)
-    assert status == _CheckStatus({"hexanator/0"}, {"hexanator/0"})
+    assert status == _CheckStatus({"hexanator/0"}, {"hexanator/0"}, set())
 
 
 @pytest.fixture

tests/unit/test_idle_check_subordinate.py

@@ -17,7 +17,7 @@
 
 def test_subordinate_apps(response: dict[str, Any], kwargs):
     status = _check(convert(response), **kwargs)
-    assert status == _CheckStatus({"ntp/0", "ubuntu/0"}, {"ntp/0", "ubuntu/0"})
+    assert status == _CheckStatus({"ntp/0", "ubuntu/0"}, {"ntp/0", "ubuntu/0"}, set())
 
 
 def test_subordinate_is_selective(response, kwargs):
@@ -26,7 +26,7 @@ def test_subordinate_is_selective(response, kwargs):
     ]
     subordinates["some-other/0"] = subordinates["ntp/0"]
     status = _check(convert(response), **kwargs)
-    assert status == _CheckStatus({"ntp/0", "ubuntu/0"}, {"ntp/0", "ubuntu/0"})
+    assert status == _CheckStatus({"ntp/0", "ubuntu/0"}, {"ntp/0", "ubuntu/0"}, set())
 
 
 @pytest.fixture