From 19ce5f889eae9aab60ea355ebbb06195d8f372a9 Mon Sep 17 00:00:00 2001
From: Gabriel Cocenza <gabriel.cocenza@canonical.com>
Date: Wed, 24 Apr 2024 12:10:22 -0300
Subject: [PATCH] Run update-status hook before checking workload version
 upgrade (#385)

- Sometimes the _get_reached_expected_target_step post upgrade step
fails because the application turn into active/idle and might take an
extra time to actually show the change on the workload version. Running
the update-status hook before checking can help reducing false
positives.
- Changed the error message to say to the user wait some minutes if the
workload version change check fails.

Closes #384
---
 cou/apps/base.py             | 10 +++++++++-
 tests/unit/apps/test_core.py |  5 ++++-
 2 files changed, 13 insertions(+), 2 deletions(-)

diff --git a/cou/apps/base.py b/cou/apps/base.py
index ddb0da63..3c2375ca 100644
--- a/cou/apps/base.py
+++ b/cou/apps/base.py
@@ -15,6 +15,7 @@
 """Base application class."""
 from __future__ import annotations
 
+import asyncio
 import logging
 import os
 from collections import defaultdict
@@ -350,6 +351,11 @@ async def _verify_workload_upgrade(self, target: OpenStackRelease, units: list[U
         :type units: list[Unit]
         :raises ApplicationError: When the workload version of the charm doesn't upgrade.
         """
+        # NOTE (gabrielcocenza) force the update-status hook on units
+        # to update the workload version
+        tasks = [self.model.run_on_unit(unit.name, "hooks/update-status") for unit in units]
+        await asyncio.gather(*tasks)
+
         status = await self.model.get_status()
         app_status = status.applications.get(self.name)
         units_not_upgraded = []
@@ -363,7 +369,9 @@ async def _verify_workload_upgrade(self, target: OpenStackRelease, units: list[U
 
         if units_not_upgraded:
             raise ApplicationError(
-                f"Cannot upgrade units '{', '.join(units_not_upgraded)}' to {target}."
+                f"Unit(s) '{', '.join(units_not_upgraded)}' did not complete the upgrade to "
+                f"{target}. Some local processes may still be executing; you may try re-running "
+                "COU in a few minutes."
             )
 
     def upgrade_plan_sanity_checks(
diff --git a/tests/unit/apps/test_core.py b/tests/unit/apps/test_core.py
index 76a6e43d..9257e75b 100644
--- a/tests/unit/apps/test_core.py
+++ b/tests/unit/apps/test_core.py
@@ -124,7 +124,10 @@ async def test_application_verify_workload_upgrade(model):
 async def test_application_verify_workload_upgrade_fail(model):
     """Test Kyestone application check unsuccessful upgrade."""
     target = OpenStackRelease("victoria")
-    exp_msg = "Cannot upgrade units 'keystone/0' to victoria."
+    exp_msg = (
+        r"Unit\(s\) 'keystone/0' did not complete the upgrade to victoria. Some local processes "
+        r"may still be executing; you may try re-running COU in a few minutes."
+    )
     machines = {"0": MagicMock(spec_set=Machine)}
     app = Keystone(
         name="keystone",