Merge pull request #517 from sed-i/feature/wait_for_status

#517

null

Author: jujubot

juju/model.py

@@ -8,6 +8,7 @@
 import re
 import stat
 import tempfile
+import warnings
 import weakref
 import zipfile
 from concurrent.futures import CancelledError
@@ -2227,7 +2228,8 @@ async def _get_source_api(self, url, controller_name=None):
         return controller
 
     async def wait_for_idle(self, apps=None, raise_on_error=True, raise_on_blocked=False,
-                            wait_for_active=False, timeout=10 * 60, idle_period=15, check_freq=0.5):
+                            wait_for_active=False, timeout=10 * 60, idle_period=15, check_freq=0.5,
+                            status=None):
         """Wait for applications in the model to settle into an idle state.
 
         :param apps (list[str]): Optional list of specific app names to wait on.
@@ -2259,7 +2261,14 @@ async def wait_for_idle(self, apps=None, raise_on_error=True, raise_on_blocked=F
 
         :param check_freq (float): How frequently, in seconds, to check the model.
             The default is every half-second.
+
+        :param status (str): The status to wait for. If None, not waiting.
+            The default is None (not waiting for any status).
         """
+        if wait_for_active:
+            warnings.warn("wait_for_active is deprecated; use status", DeprecationWarning)
+            status = "active"
+
         timeout = timedelta(seconds=timeout) if timeout is not None else None
         idle_period = timedelta(seconds=idle_period)
         start_time = datetime.now()
@@ -2311,8 +2320,8 @@ def _raise_for_status(entities, status):
                     if raise_on_blocked and unit.workload_status == "blocked":
                         blocks.setdefault("Unit", []).append(unit.name)
                         continue
-                    waiting_for_active = wait_for_active and unit.workload_status != "active"
-                    if not waiting_for_active and unit.agent_status == "idle":
+                    waiting_for_status = status and unit.workload_status != status
+                    if not waiting_for_status and unit.agent_status == "idle":
                         now = datetime.now()
                         idle_start = idle_times.setdefault(unit.name, now)
                         if now - idle_start < idle_period:

tests/integration/test_application.py

@@ -151,10 +151,10 @@ async def test_upgrade_local_charm(event_loop):
         tests_dir = Path(__file__).absolute().parent
         charm_path = tests_dir / 'upgrade-charm'
         app = await model.deploy('ubuntu', series='focal')
-        await model.wait_for_idle(wait_for_active=True)
+        await model.wait_for_idle(status="active")
         assert app.data['charm-url'].startswith('cs:ubuntu')
         await app.upgrade_charm(path=charm_path)
-        await model.wait_for_idle()  # nb: can't use wait_for_active because test charm goes to "waiting"
+        await model.wait_for_idle(status="waiting")
         assert app.data['charm-url'] == 'local:focal/ubuntu-0'
 
 

tests/integration/test_model.py

@@ -97,7 +97,7 @@ async def test_deploy_local_charm(event_loop):
     async with base.CleanModel() as model:
         await model.deploy(str(charm_path))
         assert 'charm' in model.applications
-        await model.wait_for_idle(wait_for_active=True)
+        await model.wait_for_idle(status="active")
         assert model.units['charm/0'].workload_status == 'active'
 
 
@@ -113,7 +113,7 @@ async def test_wait_local_charm_blocked(event_loop):
         assert 'charm' in model.applications
         await model.wait_for_idle()
         with pytest.raises(JujuUnitError):
-            await model.wait_for_idle(wait_for_active=True,
+            await model.wait_for_idle(status="active",
                                       raise_on_blocked=True,
                                       timeout=30)
 
@@ -130,7 +130,7 @@ async def test_wait_local_charm_waiting_timeout(event_loop):
         assert 'charm' in model.applications
         await model.wait_for_idle()
         with pytest.raises(asyncio.TimeoutError):
-            await model.wait_for_idle(wait_for_active=True, timeout=30)
+            await model.wait_for_idle(status="active", timeout=30)
 
 
 @base.bootstrapped

tests/unit/test_model.py

@@ -1,8 +1,11 @@
 import unittest
+from unittest.mock import patch, PropertyMock
 
 import mock
 
+import asyncio
 import asynctest
+import datetime
 
 from juju.client.jujudata import FileJujuData
 from juju.model import Model
@@ -262,3 +265,50 @@ async def test_with_posargs(self, mock_connect, mock_connect_model, _):
                                         macaroons='macaroons',
                                         loop='loop',
                                         max_frame_size='max_frame_size')
+
+
+# Patch timedelta to immediately force a timeout to avoid introducing an unnecessary delay in the test failing.
+# It should be safe to always set it up to lead to a timeout.
+@patch('juju.model.timedelta', new=lambda *a, **kw: datetime.timedelta(0))
+class TestModelWaitForIdle(asynctest.TestCase):
+    async def test_no_args(self):
+        m = Model()
+        with self.assertWarns(DeprecationWarning):
+            # no apps so should return right away
+            await m.wait_for_idle(wait_for_active=True)
+
+    async def test_timeout(self):
+        m = Model()
+        with self.assertRaises(asyncio.TimeoutError) as cm:
+            # no apps so should timeout after timeout period
+            await m.wait_for_idle(apps=["nonexisting_app"])
+        self.assertEqual(str(cm.exception), "Timed out waiting for model:\nnonexisting_app (missing)")
+
+    async def test_wait_for_active_status(self):
+        # create a custom apps mock
+        from types import SimpleNamespace
+        apps = {"dummy_app": SimpleNamespace(
+            status="active",
+            units=[SimpleNamespace(
+                name="mockunit/0",
+                workload_status="active",
+                workload_status_message="workload_status_message",
+                machine=None,
+                agent_status="idle",
+            )],
+        )}
+
+        with patch.object(Model, 'applications', new_callable=PropertyMock) as mock_apps:
+            mock_apps.return_value = apps
+            m = Model()
+
+            # pass "active" via `status` (str)
+            await m.wait_for_idle(apps=["dummy_app"], status="active")
+
+            # pass "active" via `wait_for_active` (bool; deprecated)
+            await m.wait_for_idle(apps=["dummy_app"], wait_for_active=True)
+
+            # use both `status` and `wait_for_active` - `wait_for_active` takes precedence
+            await m.wait_for_idle(apps=["dummy_app"], wait_for_active=True, status="doesn't matter")
+
+        mock_apps.assert_called_with()

tox.ini

@@ -10,6 +10,9 @@ skipsdist=True
 [pytest]
 markers =
     serial: mark a test that must run by itself
+filterwarnings =
+    ignore::DeprecationWarning:asynctest
+    ignore::DeprecationWarning:websockets
 
 [testenv]
 usedevelop=True
@@ -49,6 +52,14 @@ commands =
     pip install pylxd
     py.test --tb native -ra -v -n auto -k 'integration' -m 'not serial' {posargs}
 
+[testenv:unit]
+envdir = {toxworkdir}/py3
+commands =
+    # These need to be installed in a specific order
+    pip install urllib3==1.25.7
+    pip install pylxd
+    py.test --tb native -ra -v -n auto {toxinidir}/tests/unit {posargs}
+
 [testenv:serial]
 # tests that can't be run in parallel
 envdir = {toxworkdir}/py3