Merge pull request #905 from cderici/wait-for-idle-consider-app-status

#905

#### Description

This adds logic to `wait_for_idle` to also consider application status when computing status via the unit statuses. 

Should fix #831 and #900

#### QA Steps

This adds some scenarios to the existing unit tests for `wait_for_idle`. Also adds a pytest mark "wait_for_idle", so only the tests related to `wait_for_idle` can be run as a subset:

```
tox -e py3 -- -m wait_for_idle
```

```
tox -e integration -- -m wait_for_idle
```

#### Notes & Discussion

* As I wrote a TODO in the comments, `wait_for_idle` is becoming more and more convoluted. Basically we need two versions of this function. 1) `wait_for_applications()` for users who want to wait on an application to be in a certain state and doesn't particularly care about units, and 2) `wait_for_units()` for users who want finer detailed control and want to wait on certain number of units in certain applications to be in certain states etc.
* @DnPlas, @ca-scribner, @beliaev-maksim, I'm gonna ask your help validating this patch. We got a bunch of unit and integration tests that should be covering most of the scenarios related to the `wait_for_idle`, however, I'd greatly appreciate if you guys could also give this PR a good stress test. I'm basically asking you to run all the integration test you have that uses `wait_for_idle` against this patch and see if it can take a punch. Thanks!

Author: jujubot

juju/model.py

@@ -2572,11 +2572,24 @@ def _raise_for_status(entities, status):
                     if raise_on_blocked and unit.workload_status == "blocked":
                         blocks.setdefault("Unit", []).append(unit.name)
                         continue
-                    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
+                    # TODO (cderici): we need two versions of wait_for_idle, one for waiting on
+                    #  individual units, another one for waiting for an application.
+                    #  The convoluted logic below is the result of trying to do both at the same
+                    #  time
+                    need_to_wait_more_for_a_particular_status = status and (unit.workload_status != status)
+                    app_is_in_desired_status = (not status) or (app_status == status)
+                    if not need_to_wait_more_for_a_particular_status and \
+                            unit.agent_status == "idle" and \
+                            (wait_for_units or app_is_in_desired_status):
+                        # A unit is ready if either:
+                        # 1) Don't need to wait more for a particular status and the agent is "idle"
+                        # 2) We're looking for a particular status and the unit's workload,
+                        # as well as the application, is in that status. If the user wants to
+                        # see only a particular number of units in that state -- i.e. a subset of
+                        # the units is needed, then we don't care about the application status
+                        # (because e.g. app can be in 'waiting' while unit.0 is 'active' and unit.1
+                        # is 'waiting')
+
                         # 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)

tests/integration/test_model.py

@@ -785,12 +785,12 @@ async def test_get_machines(event_loop):
 
 @base.bootstrapped
 @pytest.mark.asyncio
+@pytest.mark.wait_for_idle
 async def test_wait_for_idle_without_units(event_loop):
     async with base.CleanModel() as model:
         await model.deploy(
             'ubuntu',
             application_name='ubuntu',
-            series='bionic',
             channel='stable',
             num_units=0,
         )
@@ -800,12 +800,12 @@ async def test_wait_for_idle_without_units(event_loop):
 
 @base.bootstrapped
 @pytest.mark.asyncio
+@pytest.mark.wait_for_idle
 async def test_wait_for_idle_with_not_enough_units(event_loop):
     async with base.CleanModel() as model:
         await model.deploy(
             'ubuntu',
             application_name='ubuntu',
-            series='bionic',
             channel='stable',
             num_units=2,
         )
@@ -815,6 +815,7 @@ async def test_wait_for_idle_with_not_enough_units(event_loop):
 
 @base.bootstrapped
 @pytest.mark.asyncio
+@pytest.mark.wait_for_idle
 async def test_wait_for_idle_more_units_than_needed(event_loop):
     async with base.CleanModel() as model:
         charm_path = TESTS_DIR / 'charm'
@@ -837,12 +838,12 @@ async def test_wait_for_idle_more_units_than_needed(event_loop):
 
 @base.bootstrapped
 @pytest.mark.asyncio
+@pytest.mark.wait_for_idle
 async def test_wait_for_idle_with_enough_units(event_loop):
     async with base.CleanModel() as model:
         await model.deploy(
             'ubuntu',
             application_name='ubuntu',
-            series='bionic',
             channel='stable',
             num_units=3,
         )
@@ -851,12 +852,12 @@ async def test_wait_for_idle_with_enough_units(event_loop):
 
 @base.bootstrapped
 @pytest.mark.asyncio
+@pytest.mark.wait_for_idle
 async def test_wait_for_idle_with_exact_units(event_loop):
     async with base.CleanModel() as model:
         await model.deploy(
             'ubuntu',
             application_name='ubuntu',
-            series='bionic',
             channel='stable',
             num_units=2,
         )
@@ -865,6 +866,7 @@ async def test_wait_for_idle_with_exact_units(event_loop):
 
 @base.bootstrapped
 @pytest.mark.asyncio
+@pytest.mark.wait_for_idle
 async def test_wait_for_idle_with_exact_units_scale_down(event_loop):
     """Deploys 3 units, waits for them to be idle, then removes 2 of them,
     then waits for exactly 1 unit to be left.
@@ -874,7 +876,7 @@ async def test_wait_for_idle_with_exact_units_scale_down(event_loop):
         app = await model.deploy(
             'ubuntu',
             application_name='ubuntu',
-            series='bionic',
+            series='jammy',
             channel='stable',
             num_units=3,
         )

tests/unit/test_model.py

@@ -296,25 +296,23 @@ async def test_timeout(self):
         self.assertEqual(str(cm.exception), "Timed out waiting for model:\nnonexisting_app (missing)")
 
     @pytest.mark.asyncio
+    @pytest.mark.wait_for_idle
     async def test_wait_for_active_status(self):
+        app_status = 'active'
         # create a custom apps mock
         from types import SimpleNamespace
         app = SimpleNamespace(
-            status="active",
+            status=app_status,
             units=[SimpleNamespace(
                 name="mockunit/0",
-                workload_status="active",
+                workload_status='active',
                 workload_status_message="workload_status_message",
                 machine=None,
                 agent_status="idle",
             )],
         )
-        # This is a small hack to act like we're getting 'unknown'
-        # from the api (the get_status() call), which shouldn't
-        # change the semantics of this test, as the 'unknown'
-        # has the lowest severity (so the app's 'active' status
-        # will overrule it)
-        app.get_status = base.AsyncMock(return_value='unknown')
+
+        app.get_status = base.AsyncMock(return_value=app_status)
         apps = {"dummy_app": app}
 
         with patch.object(Model, 'applications', new_callable=PropertyMock) as mock_apps:
@@ -331,3 +329,78 @@ async def test_wait_for_active_status(self):
             await m.wait_for_idle(apps=["dummy_app"], wait_for_active=True, status="doesn't matter")
 
         mock_apps.assert_called_with()
+
+    @pytest.mark.asyncio
+    @pytest.mark.wait_for_idle
+    async def test_wait_for_active_units_waiting_application(self):
+        # If the app is in waiting state, then wait more even if the units are ready
+        app_status = 'waiting'
+        # create a custom apps mock
+        from types import SimpleNamespace
+        app = SimpleNamespace(
+            status=app_status,
+            units=[SimpleNamespace(
+                name="mockunit/0",
+                workload_status='active',
+                workload_status_message="workload_status_message",
+                machine=None,
+                agent_status="idle",
+            ),
+                SimpleNamespace(
+                    name="mockunit/1",
+                    workload_status='active',
+                    workload_status_message="workload_status_message",
+                    machine=None,
+                    agent_status="idle",
+            )],
+        )
+
+        app.get_status = base.AsyncMock(return_value=app_status)
+        apps = {"dummy_app": app}
+
+        with patch.object(Model, 'applications', new_callable=PropertyMock) as mock_apps:
+            mock_apps.return_value = apps
+            m = Model()
+
+            with self.assertRaises(jasyncio.TimeoutError):
+                await m.wait_for_idle(apps=["dummy_app"], status="active")
+
+        mock_apps.assert_called_with()
+
+    @pytest.mark.asyncio
+    @pytest.mark.wait_for_idle
+    async def test_wait_for_active_units_waiting_for_units(self):
+        # If user wants to see a particular number of units, then application may be in a waiting
+        # state, return when there's at least that number of units in the desired state
+        app_status = 'waiting'
+        # create a custom apps mock
+        from types import SimpleNamespace
+        app = SimpleNamespace(
+            status=app_status,
+            units=[SimpleNamespace(
+                name="mockunit/0",
+                workload_status='active',
+                workload_status_message="workload_status_message",
+                machine=None,
+                agent_status="idle",
+            ),
+                SimpleNamespace(
+                    name="mockunit/1",
+                    workload_status='waiting',
+                    workload_status_message="workload_status_message",
+                    machine=None,
+                    agent_status="idle",
+            )],
+        )
+
+        app.get_status = base.AsyncMock(return_value=app_status)
+        apps = {"dummy_app": app}
+
+        with patch.object(Model, 'applications', new_callable=PropertyMock) as mock_apps:
+            mock_apps.return_value = apps
+            m = Model()
+
+            await m.wait_for_idle(apps=["dummy_app"], status="active", wait_for_units=1,
+                                  timeout=None)
+
+        mock_apps.assert_called_with()