Test fixes

Author: dragomirp

tests/integration/ha_tests/test_async_replication.py

@@ -308,8 +308,8 @@ async def test_data_integrator_creds_keep_on_working(
     database = data_integrator_credentials["postgresql"]["database"]
 
     any_unit = second_model.applications[DATABASE_APP_NAME].units[0].name
-    primary = await get_primary(ops_test, any_unit)
-    address = get_unit_address(ops_test, primary)
+    primary = await get_primary(ops_test, any_unit, second_model)
+    address = second_model.units.get(primary).public_address
 
     connstr = f"dbname='{database}' user='{user}' host='{address}' port='5432' password='{password}' connect_timeout=1"
     try:
@@ -319,12 +319,13 @@ async def test_data_integrator_creds_keep_on_working(
         connection.close()
 
     logger.info("Re-enable oversee users")
-    action = await primary.run_action(action_name="reenable-oversee-users")
+    leader_unit = await get_leader_unit(ops_test, DATABASE_APP_NAME, model=second_model)
+    action = await leader_unit.run_action(action_name="reenable-oversee-users")
     await action.wait()
 
     async with ops_test.fast_forward():
         await sleep(20)
-    second_model.wait_for_idle(
+    await second_model.wait_for_idle(
         apps=[DATABASE_APP_NAME],
         status="active",
         timeout=TIMEOUT,

tests/integration/helpers.py

@@ -621,17 +621,20 @@ async def get_password(ops_test: OpsTest, unit_name: str, username: str = "opera
     stop=stop_after_attempt(10),
     wait=wait_exponential(multiplier=1, min=2, max=30),
 )
-async def get_primary(ops_test: OpsTest, unit_name: str) -> str:
+async def get_primary(ops_test: OpsTest, unit_name: str, model=None) -> str:
     """Get the primary unit.
 
     Args:
         ops_test: ops_test instance.
         unit_name: the name of the unit.
+        model: Model to use.
 
     Returns:
         the current primary unit.
     """
-    action = await ops_test.model.units.get(unit_name).run_action("get-primary")
+    if not model:
+        model = ops_test.model
+    action = await model.units.get(unit_name).run_action("get-primary")
     action = await action.wait()
     return action.results["primary"]