From a74da7463e8d2d9ba619aaf4bec7fbf3519f5572 Mon Sep 17 00:00:00 2001
From: Mia Altieri <32723809+MiaAltieri@users.noreply.github.com>
Date: Tue, 14 May 2024 15:09:21 +0200
Subject: [PATCH] fix flakey sharding tests by using an accurate wait (#414)

## Issue
Sharding tests are flakey for status checking

## Solution
Wait for the model to settle before performing status checking
---
 tests/integration/sharding_tests/test_sharding.py         | 2 +-
 .../integration/sharding_tests/test_sharding_relations.py | 8 ++++++++
 2 files changed, 9 insertions(+), 1 deletion(-)

diff --git a/tests/integration/sharding_tests/test_sharding.py b/tests/integration/sharding_tests/test_sharding.py
index 462ea8e1c..6f1003ad3 100644
--- a/tests/integration/sharding_tests/test_sharding.py
+++ b/tests/integration/sharding_tests/test_sharding.py
@@ -75,7 +75,7 @@ async def test_build_and_deploy(ops_test: OpsTest) -> None:
             timeout=TIMEOUT,
         ),
         ops_test.model.wait_for_idle(
-            apps=[SHARD_ONE_APP_NAME, SHARD_TWO_APP_NAME],
+            apps=[SHARD_ONE_APP_NAME, SHARD_TWO_APP_NAME, SHARD_THREE_APP_NAME],
             status="blocked",
             idle_period=20,
             timeout=TIMEOUT,
diff --git a/tests/integration/sharding_tests/test_sharding_relations.py b/tests/integration/sharding_tests/test_sharding_relations.py
index 0338623fb..d4f9c9661 100644
--- a/tests/integration/sharding_tests/test_sharding_relations.py
+++ b/tests/integration/sharding_tests/test_sharding_relations.py
@@ -136,6 +136,7 @@ async def test_cannot_use_db_relation(ops_test: OpsTest) -> None:
 
     await ops_test.model.wait_for_idle(
         apps=SHARDING_COMPONENTS,
+        status="blocked",
         idle_period=20,
         raise_on_blocked=False,
         timeout=TIMEOUT,
@@ -172,6 +173,7 @@ async def test_cannot_use_legacy_db_relation(ops_test: OpsTest) -> None:
 
     await ops_test.model.wait_for_idle(
         apps=SHARDING_COMPONENTS,
+        status="blocked",
         idle_period=20,
         raise_on_blocked=False,
         timeout=TIMEOUT,
@@ -211,6 +213,7 @@ async def test_replication_config_server_relation(ops_test: OpsTest):
 
     await ops_test.model.wait_for_idle(
         apps=[REPLICATION_APP_NAME],
+        status="blocked",
         idle_period=20,
         raise_on_blocked=False,
         timeout=TIMEOUT,
@@ -240,6 +243,7 @@ async def test_replication_shard_relation(ops_test: OpsTest):
 
     await ops_test.model.wait_for_idle(
         apps=[REPLICATION_APP_NAME],
+        status="blocked",
         idle_period=20,
         raise_on_blocked=False,
         timeout=TIMEOUT,
@@ -277,6 +281,7 @@ async def test_replication_mongos_relation(ops_test: OpsTest) -> None:
     await ops_test.model.wait_for_idle(
         apps=[REPLICATION_APP_NAME],
         idle_period=20,
+        status="blocked",
         raise_on_blocked=False,
         timeout=TIMEOUT,
     )
@@ -313,6 +318,7 @@ async def test_shard_mongos_relation(ops_test: OpsTest) -> None:
 
     await ops_test.model.wait_for_idle(
         apps=[SHARD_ONE_APP_NAME],
+        status="blocked",
         idle_period=20,
         raise_on_blocked=False,
         timeout=TIMEOUT,
@@ -332,6 +338,7 @@ async def test_shard_mongos_relation(ops_test: OpsTest) -> None:
 
     await ops_test.model.wait_for_idle(
         apps=[SHARD_ONE_APP_NAME],
+        status="blocked",
         idle_period=20,
         raise_on_blocked=False,
         timeout=TIMEOUT,
@@ -351,6 +358,7 @@ async def test_shard_s3_relation(ops_test: OpsTest) -> None:
     await ops_test.model.wait_for_idle(
         apps=[SHARD_ONE_APP_NAME],
         idle_period=20,
+        status="blocked",
         raise_on_blocked=False,
         timeout=TIMEOUT,
     )