From 2aa46343c613a187e6df79af25bdcd250a71354c Mon Sep 17 00:00:00 2001
From: Dima Tisnek <dima.tisnek@canonical.com>
Date: Tue, 3 Dec 2024 21:09:35 +0900
Subject: [PATCH] chore: implement model.wait_for_idle(timeout=n)

---
 juju/model.py                    | 15 +++++++--------
 tests/unit/test_wait_for_idle.py |  3 ---
 2 files changed, 7 insertions(+), 11 deletions(-)

diff --git a/juju/model.py b/juju/model.py
index 47168c56b..19da4305f 100644
--- a/juju/model.py
+++ b/juju/model.py
@@ -12,6 +12,7 @@
 import stat
 import sys
 import tempfile
+import time
 import warnings
 import weakref
 import zipfile
@@ -3281,11 +3282,15 @@ async def new_wait_for_idle(
         apps = apps or list(self.applications)
         idle_times: dict[str, datetime] = {}
         units_ready: set[str] = set()  # The units that are in the desired state
-        last_log_time: list[datetime | None] = [None]
 
-        start_time = datetime.now()
+        deadline = None if timeout is None else time.monotonic() + timeout
 
         while True:
+            if deadline and time.monotonic() > deadline:
+                raise jasyncio.TimeoutError(
+                    "Timed out waiting for model to become idle"
+                )
+
             if await self._check_idle(
                 apps=apps,
                 raise_on_error=raise_on_error,
@@ -3293,13 +3298,10 @@ async def new_wait_for_idle(
                 status=status,
                 wait_for_at_least_units=wait_for_at_least_units,
                 wait_for_exact_units=wait_for_exact_units,
-                timeout=timeout,
                 idle_period=idle_period,
                 _wait_for_units=_wait_for_units,
                 idle_times=idle_times,
                 units_ready=units_ready,
-                last_log_time=last_log_time,
-                start_time=start_time,
             ):
                 break
 
@@ -3314,13 +3316,10 @@ async def _check_idle(
         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)
diff --git a/tests/unit/test_wait_for_idle.py b/tests/unit/test_wait_for_idle.py
index 09496743e..28d4953ac 100644
--- a/tests/unit/test_wait_for_idle.py
+++ b/tests/unit/test_wait_for_idle.py
@@ -327,13 +327,10 @@ def kwargs() -> dict[str, Any]:
         status=None,
         wait_for_at_least_units=None,
         wait_for_exact_units=None,
-        timeout=100,
         idle_period=0,
         _wait_for_units=1,
         idle_times={},
         units_ready=set(),
-        last_log_time=[None],
-        start_time=datetime.now(),
     )