Improve integration test

Author: marceloneppel

tests/integration/helpers.py

@@ -2,6 +2,7 @@
 # Copyright 2022 Canonical Ltd.
 # See LICENSE file for licensing details.
 import itertools
+from datetime import datetime
 from pathlib import Path
 from typing import List
 
@@ -90,20 +91,23 @@ async def check_database_creation(ops_test: OpsTest, database: str) -> None:
     stop=stop_after_attempt(10),
     wait=wait_exponential(multiplier=1, min=2, max=30),
 )
-async def check_patroni(ops_test: OpsTest, unit_name: str) -> bool:
+async def check_patroni(ops_test: OpsTest, unit_name: str, restart_time: float) -> bool:
     """Check if Patroni is running correctly on a specific unit.
 
     Args:
         ops_test: The ops test framework instance
         unit_name: The name of the unit
+        restart_time: Point in time before the unit was restarted.
 
     Returns:
         whether Patroni is running correctly.
     """
     unit_ip = await get_unit_address(ops_test, unit_name)
-    health_info = requests.get(f"http://{unit_ip}:8008/health")
-    print(health_info.json())
-    return health_info.json()["state"] == "running"
+    health_info = requests.get(f"http://{unit_ip}:8008/health").json()
+    postmaster_start_time = datetime.strptime(
+        health_info["postmaster_start_time"], "%Y-%m-%d %H:%M:%S.%f%z"
+    ).timestamp()
+    return postmaster_start_time > restart_time and health_info["state"] == "running"
 
 
 def convert_records_to_dict(records: List[tuple]) -> dict:
@@ -226,22 +230,6 @@ async def get_password(ops_test: OpsTest, username: str = "operator"):
     return result.results[f"{username}-password"]
 
 
-async def get_postgresql_start_time(ops_test: OpsTest, unit_name: str) -> bool:
-    """Get PostgreSQL start time.
-
-    Args:
-        ops_test: The ops test framework instance
-        unit_name: The name of the unit
-
-    Returns:
-        PostgreSQL start time.
-    """
-    unit_ip = await get_unit_address(ops_test, unit_name)
-    health_info = requests.get(f"http://{unit_ip}:8008/health")
-    print(health_info.json())
-    return health_info.json()["postmaster_start_time"]
-
-
 async def get_primary(ops_test: OpsTest, unit_id=0) -> str:
     """Get the primary unit.
 

tests/integration/test_password_rotation.py

@@ -1,6 +1,7 @@
 #!/usr/bin/env python3
 # Copyright 2022 Canonical Ltd.
 # See LICENSE file for licensing details.
+import time
 
 import pytest
 from pytest_operator.plugin import OpsTest
@@ -64,14 +65,8 @@ async def test_password_rotation(ops_test: OpsTest):
 
     # Restart Patroni on any non-leader unit and check that
     # Patroni and PostgreSQL continue to work.
-    non_leader_units = [
-        unit.name
-        for unit in ops_test.model.applications[APP_NAME].units
-        if not await unit.is_leader_from_status()
-    ]
-
-    for unit in non_leader_units:
-        await restart_patroni(ops_test, unit)
-
-    for unit in non_leader_units:
-        assert await check_patroni(ops_test, unit)
+    restart_time = time.time()
+    for unit in ops_test.model.applications[APP_NAME].units:
+        if not await unit.is_leader_from_status():
+            await restart_patroni(ops_test, unit.name)
+            assert await check_patroni(ops_test, unit.name, restart_time)