[DPE-5666] Extend wait_until to also wait on cos-agent (#478)

Add `cos-agent` to the wait list.

Closes #475

---------

Co-authored-by: Mehdi Bendriss <[email protected]>

tests/integration/plugins/test_plugins.py

@@ -82,7 +82,11 @@ async def _set_config(ops_test: OpsTest, deploy_type: str, conf: dict[str, str])
     await ops_test.model.applications[APP_NAME].set_config(conf)
 
 
-async def _wait_for_units(ops_test: OpsTest, deployment_type: str) -> None:
+async def _wait_for_units(
+    ops_test: OpsTest,
+    deployment_type: str,
+    wait_for_cos: bool = False,
+) -> None:
     """Wait for all units to be active.
 
     This wait will behavior accordingly to small/large.
@@ -93,10 +97,18 @@ async def _wait_for_units(ops_test: OpsTest, deployment_type: str) -> None:
             apps=[APP_NAME],
             apps_statuses=["active"],
             units_statuses=["active"],
-            wait_for_exact_units={APP_NAME: 3},
             timeout=1800,
+            wait_for_exact_units={APP_NAME: 3},
             idle_period=IDLE_PERIOD,
         )
+        if wait_for_cos:
+            await wait_until(
+                ops_test,
+                apps=[COS_APP_NAME],
+                units_statuses=["blocked"],
+                timeout=1800,
+                idle_period=IDLE_PERIOD,
+            )
         return
     await wait_until(
         ops_test,
@@ -106,17 +118,25 @@ async def _wait_for_units(ops_test: OpsTest, deployment_type: str) -> None:
             FAILOVER_ORCHESTRATOR_NAME,
             APP_NAME,
         ],
-        apps_statuses=["active"],
-        units_statuses=["active"],
         wait_for_exact_units={
             TLS_CERTIFICATES_APP_NAME: 1,
             MAIN_ORCHESTRATOR_NAME: 1,
             FAILOVER_ORCHESTRATOR_NAME: 2,
             APP_NAME: 1,
         },
+        apps_statuses=["active"],
+        units_statuses=["active"],
         timeout=1800,
         idle_period=IDLE_PERIOD,
     )
+    if wait_for_cos:
+        await wait_until(
+            ops_test,
+            apps=[COS_APP_NAME],
+            units_statuses=["blocked"],
+            timeout=1800,
+            idle_period=IDLE_PERIOD,
+        )
 
 
 @pytest.mark.parametrize("deploy_type", SMALL_DEPLOYMENTS)
@@ -173,7 +193,7 @@ async def test_prometheus_exporter_enabled_by_default(ops_test, deploy_type: str
 async def test_small_deployments_prometheus_exporter_cos_relation(ops_test, deploy_type: str):
     await ops_test.model.deploy(COS_APP_NAME, channel="edge", series=SERIES),
     await ops_test.model.integrate(APP_NAME, COS_APP_NAME)
-    await _wait_for_units(ops_test, deploy_type)
+    await _wait_for_units(ops_test, deploy_type, wait_for_cos=True)
 
     # Check that the correct settings were successfully communicated to grafana-agent
     cos_leader_id = await get_leader_unit_id(ops_test, COS_APP_NAME)
@@ -263,7 +283,7 @@ async def test_large_deployment_prometheus_exporter_cos_relation(ops_test, deplo
     await ops_test.model.integrate(MAIN_ORCHESTRATOR_NAME, COS_APP_NAME)
     await ops_test.model.integrate(APP_NAME, COS_APP_NAME)
 
-    await _wait_for_units(ops_test, deploy_type)
+    await _wait_for_units(ops_test, deploy_type, wait_for_cos=True)
 
     leader_id = await get_leader_unit_id(ops_test, APP_NAME)
     leader_name = f"{APP_NAME}/{leader_id}"
@@ -315,7 +335,7 @@ async def test_prometheus_monitor_user_password_change(ops_test, deploy_type: st
     result1 = await run_action(
         ops_test, leader_id, "set-password", {"username": "monitor"}, app=app
     )
-    await _wait_for_units(ops_test, deploy_type)
+    await _wait_for_units(ops_test, deploy_type, wait_for_cos=True)
 
     new_password = result1.response.get("monitor-password")
     # Now, we compare the change in the action above with the opensearch's nodes.