diff --git a/juju/model.py b/juju/model.py
index 746a5208..f1e0e2e1 100644
--- a/juju/model.py
+++ b/juju/model.py
@@ -2870,15 +2870,15 @@ async def _get_source_api(self, url):
     # - may be used with N apps, N units with a mapping {app_name: unit_count}
     async def wait_for_idle(self,
         apps: Optional[List[str]] = None,
-        raise_on_error=True,
-        raise_on_blocked=False,
-        wait_for_active=False,
-        timeout=10 * 60,
-        idle_period=15,
+        raise_on_error: bool = True,
+        raise_on_blocked: bool = False,
+        wait_for_active: bool = False,
+        timeout: Optional[float] = 10 * 60,
+        idle_period: float = 15,
         check_freq=0.5,
-        status=None,
-        wait_for_at_least_units=None,
-        wait_for_exact_units: Union[None, int, Dict[str, int]] = None,
+        status: Optional[str] = None,
+        wait_for_at_least_units: Optional[int] = None,
+        wait_for_exact_units: Optional[int] = None,
     ):
         """Wait for applications in the model to settle into an idle state.
 
@@ -2930,8 +2930,6 @@ async def wait_for_idle(self,
 
         _wait_for_units = wait_for_at_least_units if wait_for_at_least_units is not None else 1
 
-        timeout = timedelta(seconds=timeout) if timeout is not None else None
-        idle_period = timedelta(seconds=idle_period)
         start_time = datetime.now()
         # Type check against the common error of passing a str for apps
         if apps is not None and (not isinstance(apps, list) or
@@ -2946,15 +2944,22 @@ async def wait_for_idle(self,
         units_ready: Set[str] = set()  # The units that are in the desired state
         last_log_time: List[Optional[datetime]] = [None]
 
+        # FIXME seems inconsistent wrt. use in integration tests
         if wait_for_exact_units is not None:
             assert isinstance(wait_for_exact_units, int) and wait_for_exact_units >= 0, \
                 'Invalid value for wait_for_exact_units : %s' % wait_for_exact_units
 
         while True:
-            if (await self._wait_for_idle_cycle(
+            if (await self._check_idle(
                 apps=apps,
                 raise_on_error=raise_on_error,
                 raise_on_blocked=raise_on_blocked,
+                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,
@@ -2964,17 +2969,25 @@ async def wait_for_idle(self,
 
             await jasyncio.sleep(check_freq)
 
-    async def _wait_for_idle_cycle(
+    async def _check_idle(
         self,
         *,
         apps: List[str],
         raise_on_error: bool,
         raise_on_blocked: bool,
+        status: Optional[str],
+        wait_for_at_least_units: Optional[int],
+        wait_for_exact_units: Optional[int],
+        timeout: Optional[float],
+        idle_period: float,
+        _wait_for_units: int,
         idle_times: Dict[str, datetime] = {},
         units_ready: Set[str] = set(),  # The units that are in the desired state
         last_log_time: List[Optional[datetime]] = [None],
         start_time: datetime = datetime.now(),
     ):
+        _timeout = timedelta(seconds=timeout) if timeout is not None else None
+        _idle_period = timedelta(seconds=idle_period)
         log_interval = timedelta(seconds=30)
         # The list 'busy' is what keeps this loop going,
         # i.e. it'll stop when busy is empty after all the
@@ -3028,6 +3041,7 @@ async def _wait_for_idle_cycle(
                             (len(app.units), _wait_for_units))
                 return False
             # User is waiting for at least a certain # of units, and we have enough
+            # FIXME seems to skip unit status checks, is that expected?
             elif wait_for_at_least_units and len(units_ready) >= _wait_for_units:
                 # So no need to keep looking, we have the desired number of units ready to go,
                 # exit the loop. Don't just return here, though, we might still have some
@@ -3085,7 +3099,7 @@ async def _wait_for_idle_cycle(
                     # FIXME key in units dict
                     idle_start = idle_times.setdefault(unit.name, now)
 
-                    if now - idle_start < idle_period:
+                    if now - idle_start < _idle_period:
                         busy.append("{} [{}] {}: {}".format(unit.name,
                                                             unit.agent_status,
                                                             unit.workload_status,
@@ -3104,7 +3118,7 @@ async def _wait_for_idle_cycle(
         if not busy:
             return True
 
-        if timeout is not None and datetime.now() - start_time > timeout:
+        if _timeout is not None and datetime.now() - start_time > _timeout:
             raise jasyncio.TimeoutError("\n  ".join(["Timed out waiting for model:", *busy]))
 
         if last_log_time[0] is None or datetime.now() - last_log_time[0] > log_interval: