From 59989f51e09b0b3bbe544f5fe5cdb5315917f40a Mon Sep 17 00:00:00 2001
From: Artem Yevsiukov <a.yevsiukov@gmail.com>
Date: Mon, 15 Apr 2024 14:50:04 +0300
Subject: [PATCH] added `check_tls_replication` for checking replicas encrypted
 connection (#437)

---
 tests/integration/helpers.py  | 29 +++++++++++++++++++++++++++++
 tests/integration/test_tls.py |  4 ++++
 2 files changed, 33 insertions(+)

diff --git a/tests/integration/helpers.py b/tests/integration/helpers.py
index 1845e69778..786cb67685 100644
--- a/tests/integration/helpers.py
+++ b/tests/integration/helpers.py
@@ -736,6 +736,35 @@ async def check_tls(ops_test: OpsTest, unit_name: str, enabled: bool) -> bool:
         return False
 
 
+async def check_tls_replication(ops_test: OpsTest, unit_name: str, enabled: bool) -> bool:
+    """Returns whether TLS is enabled on the replica PostgreSQL instance.
+
+    Args:
+        ops_test: The ops test framework instance.
+        unit_name: The name of the replica of the PostgreSQL instance.
+        enabled: check if TLS is enabled/disabled
+
+    Returns:
+        Whether TLS is enabled/disabled.
+    """
+    unit_address = get_unit_address(ops_test, unit_name)
+    password = await get_password(ops_test, unit_name)
+
+    # Check for the all replicas using encrypted connection
+    output = await execute_query_on_unit(
+        unit_address,
+        password,
+        "SELECT pg_ssl.ssl, pg_sa.client_addr FROM pg_stat_ssl pg_ssl"
+        " JOIN pg_stat_activity pg_sa ON pg_ssl.pid = pg_sa.pid"
+        " AND pg_sa.usename = 'replication';",
+    )
+
+    for i in range(0, len(output), 2):
+        if output[i] != enabled:
+            return False
+    return True
+
+
 async def check_tls_patroni_api(ops_test: OpsTest, unit_name: str, enabled: bool) -> bool:
     """Returns whether TLS is enabled on Patroni REST API.
 
diff --git a/tests/integration/test_tls.py b/tests/integration/test_tls.py
index 0f194d0c49..70c7150470 100644
--- a/tests/integration/test_tls.py
+++ b/tests/integration/test_tls.py
@@ -15,6 +15,7 @@
     change_primary_start_timeout,
     check_tls,
     check_tls_patroni_api,
+    check_tls_replication,
     db_connect,
     get_password,
     get_primary,
@@ -85,6 +86,9 @@ async def test_tls_enabled(ops_test: OpsTest) -> None:
             if unit.name != primary
         ][0]
 
+        # Check if TLS enabled for replication
+        assert await check_tls_replication(ops_test, primary, enabled=True)
+
         # Enable additional logs on the PostgreSQL instance to check TLS
         # being used in a later step and make the fail-over to happens faster.
         await ops_test.model.applications[DATABASE_APP_NAME].set_config({