+ @staticmethod
+ async def wait_for_units_idle(
+ model: Model, application: Application, timeout: float = 60
+ ):
+ """
+ Waits for the application and all its units to transition back to idle
+
+ :param: model: Model to observe
+ :param: application: The application to be observed
+ :param: timeout: Maximum time between two updates in the model
+
+ :raises: asyncio.TimeoutError when timeout reaches
+ """
+
+ ensure_units_idle = asyncio.ensure_future(
+ asyncio.wait_for(
+ JujuModelWatcher.ensure_units_idle(model, application), timeout
+ )
+ )
+ tasks = [
+ ensure_units_idle,
+ ]
+ (done, pending) = await asyncio.wait(
+ tasks, timeout=timeout, return_when=asyncio.FIRST_COMPLETED
+ )
+
+ if ensure_units_idle in pending:
+ ensure_units_idle.cancel()
+ raise TimeoutError(
+ "Application's units failed to return to idle after {} seconds".format(
+ timeout
+ )
+ )
+ if ensure_units_idle.result():
+ pass
+
+ @staticmethod
+ async def ensure_units_idle(model: Model, application: Application):
+ """
+ Waits forever until the application's units to transition back to idle
+
+ :param: model: Model to observe
+ :param: application: The application to be observed
+ """
+
+ try:
+ allwatcher = client.AllWatcherFacade.from_connection(model.connection())
+ unit_wanted_state = "executing"
+ final_state_reached = False
+
+ units = application.units
+ final_state_seen = {unit.entity_id: False for unit in units}
+ agent_state_seen = {unit.entity_id: False for unit in units}
+ workload_state = {unit.entity_id: False for unit in units}
+
+ try:
+ while not final_state_reached:
+ change = await allwatcher.Next()
+
+ # Keep checking to see if new units were added during the change
+ for unit in units:
+ if unit.entity_id not in final_state_seen:
+ final_state_seen[unit.entity_id] = False
+ agent_state_seen[unit.entity_id] = False
+ workload_state[unit.entity_id] = False
+
+ for delta in change.deltas:
+ await asyncio.sleep(0)
+ if delta.entity != units[0].entity_type:
+ continue
+
+ final_state_reached = True
+ for unit in units:
+ if delta.data["name"] == unit.entity_id:
+ status = delta.data["agent-status"]["current"]
+ workload_state[unit.entity_id] = delta.data[
+ "workload-status"
+ ]["current"]
+
+ if status == unit_wanted_state:
+ agent_state_seen[unit.entity_id] = True
+ final_state_seen[unit.entity_id] = False
+
+ if (
+ status == "idle"
+ and agent_state_seen[unit.entity_id]
+ ):
+ final_state_seen[unit.entity_id] = True
+
+ final_state_reached = (
+ final_state_reached
+ and final_state_seen[unit.entity_id]
+ and workload_state[unit.entity_id]
+ in [
+ "active",
+ "error",
+ ]
+ )
+
+ except ConnectionClosed:
+ pass
+ # This is expected to happen when the
+ # entity reaches its final state, because
+ # the model connection is closed afterwards
+ except Exception as e:
+ raise e
+