chore: refactor using plain old classes

Author: dimaqq

juju/model/__init__.py

@@ -3300,24 +3300,24 @@ async def new_wait_for_idle(
 
         started = time.monotonic()
         deadline = None if timeout is None else started + timeout
+        loop = _idle.Loop(
+            apps=apps,
+            wait_for_exact_units=wait_for_exact_units,
+            wait_for_units=wait_for_units,
+            idle_period=idle_period,
+        )
 
-        async def status_on_demand():
-            while True:
-                yield _idle.check(
+        while True:
+            done = loop.next(
+                _idle.check(
                     await self.get_status(),
                     apps=apps,
                     raise_on_error=raise_on_error,
                     raise_on_blocked=raise_on_blocked,
                     status=status,
                 )
+            )
 
-        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,
-        ):
             logger.info(f"wait_for_idle start{time.monotonic() - started:+.1f} {done=}")
             if done:
                 break

juju/model/_idle.py

@@ -7,7 +7,7 @@
 import logging
 import time
 from dataclasses import dataclass
-from typing import AbstractSet, AsyncIterable
+from typing import AbstractSet
 
 from ..client._definitions import (
     ApplicationStatus,
@@ -32,70 +32,72 @@ class CheckStatus:
     """Units with stable agent status (FIXME details)."""
 
 
-async def loop(
-    foo: AsyncIterable[CheckStatus | None],
-    *,
-    apps: AbstractSet[str],
-    wait_for_exact_units: int | None = None,
-    wait_for_units: int,
-    idle_period: float,
-) -> AsyncIterable[bool]:
-    """The outer, time-dependents logic of a wait_for_idle loop."""
-    idle_since: dict[str, float] = {}
-
-    async for status in foo:
+class Loop:
+    def __init__(
+        self,
+        *,
+        apps: AbstractSet[str],
+        wait_for_exact_units: int | None = None,
+        wait_for_units: int,
+        idle_period: float,
+    ):
+        self.apps = apps
+        self.wait_for_exact_units = wait_for_exact_units
+        self.wait_for_units = wait_for_units
+        self.idle_period = idle_period
+        self.idle_since: dict[str, float] = {}
+
+    def next(self, status: CheckStatus | None) -> bool:
         logger.info("wait_for_idle iteration %s", status)
         now = time.monotonic()
 
         if not status:
-            yield False
-            continue
+            return False
 
-        expected_idle_since = now - idle_period
+        expected_idle_since = now - self.idle_period
 
         # FIXME there's some confusion about what a "busy" unit is
         # 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.idle_units:
-                idle_since[name] = min(now, idle_since.get(name, float("inf")))
+                self.idle_since[name] = min(
+                    now, self.idle_since.get(name, float("inf"))
+                )
             else:
-                idle_since[name] = float("inf")
+                self.idle_since[name] = float("inf")
 
-        if busy := {n for n, t in idle_since.items() if t > expected_idle_since}:
+        if busy := {n for n, t in self.idle_since.items() if t > expected_idle_since}:
             logger.info("Waiting for units to be idle enough: %s", busy)
-            yield False
-            continue
+            return False
 
-        for app_name in apps:
+        for app_name in self.apps:
             ready_units = [
                 n for n in status.ready_units if n.startswith(f"{app_name}/")
             ]
-            if len(ready_units) < wait_for_units:
+            if len(ready_units) < self.wait_for_units:
                 logger.info(
                     "Waiting for app %r units %s >= %s",
                     app_name,
                     len(status.ready_units),
-                    wait_for_units,
+                    self.wait_for_units,
                 )
-                yield False
-                break
+                return False
 
             if (
-                wait_for_exact_units is not None
-                and len(ready_units) != wait_for_exact_units
+                self.wait_for_exact_units is not None
+                and len(ready_units) != self.wait_for_exact_units
             ):
                 logger.info(
                     "Waiting for app %r units %s == %s",
                     app_name,
                     len(ready_units),
-                    wait_for_exact_units,
+                    self.wait_for_exact_units,
                 )
-                yield False
-                break
-        else:
-            yield True
+                return False
+
+        return True
 
 
 def check(

tests/unit/test_idle_loop.py

@@ -2,50 +2,61 @@
 # Licensed under the Apache V2, see LICENCE file for details.
 from __future__ import annotations
 
-from freezegun import freeze_time
-
-from juju.model._idle import CheckStatus, loop
+from typing import AbstractSet, Iterable
 
+from freezegun import freeze_time
 
-async def alist(agen):
-    return [v async for v in agen]
+from juju.model._idle import CheckStatus, Loop
+
+
+def unroll(
+    statuses: Iterable[CheckStatus | None],
+    *,
+    apps: AbstractSet[str],
+    wait_for_exact_units: int | None = None,
+    wait_for_units: int,
+    idle_period: float,
+) -> list[bool]:
+    loop = Loop(
+        apps=apps,
+        wait_for_exact_units=wait_for_exact_units,
+        wait_for_units=wait_for_units,
+        idle_period=idle_period,
+    )
+    return [loop.next(s) for s in statuses]
 
 
-async def test_wait_for_apps():
-    async def checks():
+def test_wait_for_apps():
+    def checks():
         yield None
         yield None
 
-    assert await alist(
-        loop(
-            checks(),
-            apps={"a"},
-            wait_for_units=0,
-            idle_period=0,
-        )
+    assert unroll(
+        checks(),
+        apps={"a"},
+        wait_for_units=0,
+        idle_period=0,
     ) == [False, False]
 
 
-async def test_at_least_units():
-    async def checks():
+def test_at_least_units():
+    def checks():
         yield CheckStatus({"u/0", "u/1", "u/2"}, {"u/0"}, {"u/0", "u/1", "u/2"})
         yield CheckStatus({"u/0", "u/1", "u/2"}, {"u/0", "u/1"}, {"u/0", "u/1", "u/2"})
         yield CheckStatus(
             {"u/0", "u/1", "u/2"}, {"u/0", "u/1", "u/2"}, {"u/0", "u/1", "u/2"}
         )
 
     with freeze_time():
-        assert await alist(
-            loop(
-                checks(),
-                apps={"u"},
-                wait_for_units=2,
-                idle_period=0,
-            )
+        assert unroll(
+            checks(),
+            apps={"u"},
+            wait_for_units=2,
+            idle_period=0,
         ) == [False, True, True]
 
 
-async def test_for_exact_units():
+def test_for_exact_units():
     good = CheckStatus(
         {"u/0", "u/1", "u/2"},
         {"u/1", "u/2"},
@@ -62,55 +73,49 @@ async def test_for_exact_units():
         {"u/0", "u/1", "u/2"},
     )
 
-    async def checks():
+    def checks():
         yield too_few
         yield good
         yield too_many
         yield good
 
-    assert await alist(
-        loop(
-            checks(),
-            apps={"u"},
-            wait_for_units=1,
-            wait_for_exact_units=2,
-            idle_period=0,
-        )
+    assert unroll(
+        checks(),
+        apps={"u"},
+        wait_for_units=1,
+        wait_for_exact_units=2,
+        idle_period=0,
     ) == [False, True, False, True]
 
 
-async def test_idle_ping_pong():
+def test_idle_ping_pong():
     good = CheckStatus({"hexanator/0"}, {"hexanator/0"}, {"hexanator/0"})
     bad = CheckStatus({"hexanator/0"}, {"hexanator/0"}, set())
 
-    async def checks():
+    def checks():
         with freeze_time() as clock:
             for status in [good, bad, good, bad]:
                 yield status
                 clock.tick(10)
 
-    assert await alist(
-        loop(
-            checks(),
-            apps={"hexanator"},
-            wait_for_units=1,
-            idle_period=15,
-        )
+    assert unroll(
+        checks(),
+        apps={"hexanator"},
+        wait_for_units=1,
+        idle_period=15,
     ) == [False, False, False, False]
 
 
-async def test_idle_period():
-    async def checks():
+def test_idle_period():
+    def checks():
         with freeze_time() as clock:
             for _ in range(4):
                 yield CheckStatus({"hexanator/0"}, {"hexanator/0"}, {"hexanator/0"})
                 clock.tick(10)
 
-    assert await alist(
-        loop(
-            checks(),
-            apps={"hexanator"},
-            wait_for_units=1,
-            idle_period=15,
-        )
+    assert unroll(
+        checks(),
+        apps={"hexanator"},
+        wait_for_units=1,
+        idle_period=15,
     ) == [False, False, True, True]