Use status output for block_until_unit_wl_status

unit.workload_status was actually reporting the application workload
status. Using the full status output from model.get_status() gives us
unit by unit workload status.

Updated block_until_unit_wl_status to use the full status output to
correctly block per unit.

unit_tests/test_zaza_model.py

@@ -112,6 +112,25 @@ class TestModel(ut_utils.BaseTestCase):
         }
         self.Model_mock = mock.MagicMock()
 
+        # Juju Status Object and data
+        self.key = "instance-id"
+        self.key_data = "machine-uuid"
+        self.machine = "1"
+        self.machine_data = {self.key: self.key_data}
+        self.unit = "app/1"
+        self.unit_data = {
+            "workload-status": {"status": "active"},
+            "machine": self.machine}
+        self.application = "app"
+        self.application_data = {"units": {self.unit: self.unit_data}}
+        self.subordinate_application = "subordinate_application"
+        self.subordinate_application_data = {
+            "subordinate-to": [self.application]}
+        self.juju_status = mock.MagicMock()
+        self.juju_status.applications = {
+            self.application: self.application_data}
+        self.juju_status.machines = self.machine_data
+
         async def _connect_model(model_name):
             return model_name
 
@@ -886,34 +905,49 @@ disk_formats = ami,ari,aki,vhd,vmdk,raw,qcow2,vdi,iso,root-tar
 
     def test_block_until_unit_wl_status(self):
         async def _block_until(f, timeout=None):
-            if not f():
+            rc = await f()
+            if not rc:
                 raise asyncio.futures.TimeoutError
-        self.patch_object(model, 'get_juju_model', return_value='mname')
+
+        async def _get_status():
+            return self.juju_status
+
         self.patch_object(model, 'Model')
         self.Model.return_value = self.Model_mock
-        self.Model_mock.block_until.side_effect = _block_until
+        self.patch_object(model, 'get_juju_model', return_value='mname')
         self.patch_object(model, 'get_unit_from_name')
-        self.get_unit_from_name.return_value = mock.MagicMock(
-            workload_status='active')
+        self.patch_object(model, 'async_get_status')
+        self.async_get_status.side_effect = _get_status
+        self.patch_object(model, 'async_block_until')
+        self.async_block_until.side_effect = _block_until
         model.block_until_unit_wl_status(
-            'app/2',
+            'app/1',
             'active',
             timeout=0.1)
 
     def test_block_until_unit_wl_status_fail(self):
         async def _block_until(f, timeout=None):
-            if not f():
+            rc = await f()
+            if not rc:
                 raise asyncio.futures.TimeoutError
-        self.patch_object(model, 'get_juju_model', return_value='mname')
+
+        async def _get_status():
+            return self.juju_status
+
+        (self.juju_status.applications[self.application]
+            ["units"][self.unit]["workload-status"]["status"]) = "blocked"
+
         self.patch_object(model, 'Model')
         self.Model.return_value = self.Model_mock
-        self.Model_mock.block_until.side_effect = _block_until
+        self.patch_object(model, 'get_juju_model', return_value='mname')
         self.patch_object(model, 'get_unit_from_name')
-        self.get_unit_from_name.return_value = mock.MagicMock(
-            workload_status='maintenance')
+        self.patch_object(model, 'async_get_status')
+        self.async_get_status.side_effect = _get_status
+        self.patch_object(model, 'async_block_until')
+        self.async_block_until.side_effect = _block_until
         with self.assertRaises(asyncio.futures.TimeoutError):
             model.block_until_unit_wl_status(
-                'app/2',
+                'app/1',
                 'active',
                 timeout=0.1)
 

zaza/model.py

@@ -1090,25 +1090,27 @@ async def async_block_until_unit_wl_status(unit_name, status, model_name=None,
     blocks until the given unit has the desired workload status::
 
          block_until_unit_wl_status(
-            'modelname',
             aunit,
-            'active')
+            'active'
+            model_name='modelname')
 
-    :param model_name: Name of model to query.
-    :type model_name: str
-    :param unit_name: Name of unit to run action on
+    :param unit_name: Name of unit
     :type unit_name: str
     :param status: Status to wait for (active, maintenance etc)
     :type status: str
+    :param model_name: Name of model to query.
+    :type model_name: str
     :param timeout: Time to wait for unit to achieved desired status
     :type timeout: float
     """
-    async with run_in_model(model_name) as model:
-        unit = get_unit_from_name(unit_name, model)
-        await model.block_until(
-            lambda: unit.workload_status == status,
-            timeout=timeout
-        )
+    async def _unit_status():
+        app = unit_name.split("/")[0]
+        model_status = await async_get_status()
+        return (model_status.applications[app]['units'][unit_name]
+                ['workload-status']['status'] == status)
+
+    async with run_in_model(model_name):
+        await async_block_until(_unit_status, timeout=timeout)
 
 block_until_unit_wl_status = sync_wrapper(
     async_block_until_unit_wl_status)