From 14c14c0245ebe81607e61ea38f4ef30cf12740fd Mon Sep 17 00:00:00 2001
From: Corey Bryant <corey.bryant@canonical.com>
Date: Tue, 17 Oct 2023 19:46:24 +0000
Subject: [PATCH] Add (async_)block_until_charm_channel support

This adds async_block_until_charm_channel and
block_until_charm_channel which will initially be
used for charm upgrade testing.
---
 unit_tests/test_zaza_model.py | 27 +++++++++++++++++++++++++++
 zaza/model.py                 | 28 ++++++++++++++++++++++++++++
 2 files changed, 55 insertions(+)

diff --git a/unit_tests/test_zaza_model.py b/unit_tests/test_zaza_model.py
index 974519ace..b000fd153 100644
--- a/unit_tests/test_zaza_model.py
+++ b/unit_tests/test_zaza_model.py
@@ -29,6 +29,7 @@
     AsyncTimeoutError = asyncio.futures.TimeoutError
 
 import copy
+import collections
 import concurrent
 import datetime
 import mock
@@ -38,6 +39,7 @@
 import unit_tests.utils as ut_utils
 
 import zaza.model as model
+import zaza.utilities.ro_types as ro_types
 import zaza
 
 
@@ -1771,6 +1773,31 @@ async def _get_status(*args):
         with self.assertRaises(AsyncTimeoutError):
             model.block_until_charm_url('app', 'something wrong', timeout=0.1)
 
+    def test_block_until_charm_channel(self):
+
+        async def _block_until(f, timeout=None):
+            rc = await f()
+            if not rc:
+                raise AsyncTimeoutError
+
+        async def _get_status(*args):
+            return self.juju_status
+        self.patch_object(model, 'Model')
+        self.Model.return_value = self.Model_mock
+        self.patch_object(model, 'async_block_until')
+        self.async_block_until.side_effect = _block_until
+        self.patch_object(model, 'async_get_status')
+        self.async_get_status.side_effect = _get_status
+        target_channel = '2023.2/stable'
+        charm_channel = collections.OrderedDict(
+            {'charm_channel': target_channel})
+        self.juju_status.applications[self.application] = (
+            ro_types.resolve_immutable(charm_channel))
+        model.block_until_charm_channel('app', target_channel)
+        with self.assertRaises(AsyncTimeoutError):
+            model.block_until_charm_channel(
+                'app', 'something wrong', timeout=0.1)
+
     def block_until_service_status_base(self, rou_return):
 
         async def _block_until(f, timeout=None):
diff --git a/zaza/model.py b/zaza/model.py
index f4bd3d5df..ccec8150f 100644
--- a/zaza/model.py
+++ b/zaza/model.py
@@ -1838,6 +1838,34 @@ async def _check_charm_url():
 block_until_charm_url = sync_wrapper(async_block_until_charm_url)
 
 
+async def async_block_until_charm_channel(application, target_channel,
+                                          model_name=None, timeout=2700):
+    """Block until the charm channel matches target_channel.
+
+    An example accessing this function via its sync wrapper::
+
+        block_until_charm_channel('cinder', '2023.2/stable')
+
+    :param application_name: Name of application
+    :type application_name: str
+    :param target_channel: Target charm channel
+    :type target_channel: str
+    :param model_name: Name of model to interact with.
+    :type model_name: str
+    :param timeout: Time to wait for status to be achieved
+    :type timeout: float
+    """
+    async def _check_charm_channel():
+        model_status = await async_get_status(model_name)
+        charm_channel = model_status.applications[application].charm_channel
+        return charm_channel == target_channel
+
+    await async_block_until(_check_charm_channel, timeout=timeout)
+
+
+block_until_charm_channel = sync_wrapper(async_block_until_charm_channel)
+
+
 async def async_block_until_service_status(unit_name, services, target_status,
                                            model_name=None, timeout=2700,
                                            pgrep_full=False):