[JUJU-3583] wait_for_idle to not block when enough units are ready (#840)

* Fix wait_for_units flag to not block when enough units ready

wait_for_idle will keep waiting if there are less number of units
available than requested (via the wait_for_units flag). However, if
there are already a number of units in a desired status ready to go,
more than (or equal to) wait_for_units, then it shouldn't block until
other not-yet-available units to get into that desired state as well.

fixes #837

* Add integration test for wait_for_units in wait_for_idle

Author: cderici

juju/model.py

@@ -2457,6 +2457,7 @@ async def wait_for_idle(self, apps=None, raise_on_error=True, raise_on_blocked=F
         start_time = datetime.now()
         apps = apps or self.applications
         idle_times = {}
+        units_ready = set()  # The units that are in the desired state
         last_log_time = None
         log_interval = timedelta(seconds=30)
 
@@ -2478,6 +2479,9 @@ def _raise_for_status(entities, status):
                 ))
 
         while True:
+            # The list 'busy' is what keeps this loop going,
+            # i.e. it'll stop when busy is empty after all the
+            # units are scanned
             busy = []
             errors = {}
             blocks = {}
@@ -2490,15 +2494,20 @@ def _raise_for_status(entities, status):
                     errors.setdefault("App", []).append(app.name)
                 if raise_on_blocked and app.status == "blocked":
                     blocks.setdefault("App", []).append(app.name)
+                # Check if wait_for_exact_units flag is used
                 if wait_for_exact_units > 0:
                     if len(app.units) != wait_for_exact_units:
                         busy.append(app.name + " (waiting for exactly %s units, current : %s)" %
                                     (wait_for_exact_units, len(app.units)))
                         continue
+                # If we have less # of units then required, then wait a bit more
                 elif len(app.units) < wait_for_units:
                     busy.append(app.name + " (not enough units yet - %s/%s)" %
                                 (len(app.units), wait_for_units))
                     continue
+                elif len(units_ready) >= wait_for_units:
+                    # No need to keep looking, we have the desired number of units ready to go
+                    break
                 for unit in app.units:
                     if unit.machine is not None and unit.machine.status == "error":
                         errors.setdefault("Machine", []).append(unit.machine.id)
@@ -2512,10 +2521,17 @@ def _raise_for_status(entities, status):
                     if raise_on_blocked and unit.workload_status == "blocked":
                         blocks.setdefault("Unit", []).append(unit.name)
                         continue
-                    waiting_for_status = status and unit.workload_status != status
-                    if not waiting_for_status and unit.agent_status == "idle":
+                    waiting_for_a_particular_status = status and unit.workload_status != status
+                    if not waiting_for_a_particular_status and unit.agent_status == "idle":
+                        # We'll be here in two cases:
+                        # 1) We're not waiting for a particular status and the agent is "idle"
+                        # 2) We're waiting for a particular status and the workload is in that status
+                        # Either way, the unit is ready, start measuring the time period that
+                        # it needs to stay in that state (i.e. idle_period)
+                        units_ready.add(unit.name)
                         now = datetime.now()
                         idle_start = idle_times.setdefault(unit.name, now)
+                        print(f'unit {unit.name} is waiting since : {idle_start} -- now : {now} -- waiting for : {now - idle_start}')
                         if now - idle_start < idle_period:
                             busy.append("{} [{}] {}: {}".format(unit.name,
                                                                 unit.agent_status,

tests/integration/test_model.py

@@ -774,6 +774,28 @@ async def test_wait_for_idle_with_not_enough_units(event_loop):
             await model.wait_for_idle(timeout=5 * 60, wait_for_units=3)
 
 
+@base.bootstrapped
+@pytest.mark.asyncio
+async def test_wait_for_idle_more_units_than_needed(event_loop):
+    async with base.CleanModel() as model:
+        charm_path = TESTS_DIR / 'charm'
+
+        # we add 2 units of a local charm that does nothing
+        # (i.e. can't go into active/idle)
+        await model.deploy(str(charm_path), num_units=2)
+
+        # then add 1 unit of ubuntu charm
+        await model.deploy(
+            'ubuntu',
+            application_name='ubuntu',
+            num_units=1,
+        )
+
+        # because the wait_for_units=1, wait_for_idle should return without timing out
+        # even though there are two more units that aren't active/idle
+        await model.wait_for_idle(timeout=5 * 60, wait_for_units=1, status='active')
+
+
 @base.bootstrapped
 @pytest.mark.asyncio
 async def test_wait_for_idle_with_enough_units(event_loop):