From c200694ab4a924dd3b29f299315d4c4e257aed94 Mon Sep 17 00:00:00 2001 From: Dima Tisnek Date: Thu, 19 Dec 2024 10:58:29 +0900 Subject: [PATCH] chore: refactor using plain old classes --- juju/model/__init__.py | 20 +++---- juju/model/_idle.py | 68 +++++++++++----------- tests/unit/test_idle_loop.py | 107 ++++++++++++++++++----------------- 3 files changed, 101 insertions(+), 94 deletions(-) diff --git a/juju/model/__init__.py b/juju/model/__init__.py index 951713a7..826cff88 100644 --- a/juju/model/__init__.py +++ b/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 diff --git a/juju/model/_idle.py b/juju/model/_idle.py index 7ca0cac9..f73539fa 100644 --- a/juju/model/_idle.py +++ b/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,26 +32,29 @@ 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: @@ -59,43 +62,42 @@ async def loop( # 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( diff --git a/tests/unit/test_idle_loop.py b/tests/unit/test_idle_loop.py index 2e97bd67..300c50f3 100644 --- a/tests/unit/test_idle_loop.py +++ b/tests/unit/test_idle_loop.py @@ -2,32 +2,45 @@ # 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( @@ -35,17 +48,15 @@ async def checks(): ) 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]