From f5b20ae9b763e47077e0ab3b485f463e07195865 Mon Sep 17 00:00:00 2001
From: Pedro Guimaraes <pedro.guimaraes@canonical.com>
Date: Thu, 18 Apr 2024 20:32:03 +0200
Subject: [PATCH] Fix deletion test

---
 .../relations/pgbouncer_provider/helpers.py   | 40 ++++++++++++++++---
 .../test_pgbouncer_provider.py                |  1 -
 2 files changed, 34 insertions(+), 7 deletions(-)

diff --git a/tests/integration/relations/pgbouncer_provider/helpers.py b/tests/integration/relations/pgbouncer_provider/helpers.py
index 58b8a9ca6..700fa3d61 100644
--- a/tests/integration/relations/pgbouncer_provider/helpers.py
+++ b/tests/integration/relations/pgbouncer_provider/helpers.py
@@ -125,6 +125,24 @@ async def run_sql_on_application_charm(
     return result.results
 
 
+async def is_external_connectivity_set(
+    ops_test: OpsTest,
+    application_name: str,
+    relation_name: str,
+    relation_id: str = None,
+) -> Optional[str]:
+    data = await get_application_relation_data(
+        ops_test,
+        application_name,
+        relation_name,
+        "data",
+        relation_id,
+    )
+    if not data:
+        return False
+    return json.loads(data).get("external-node-connectivity", None) == "true"
+
+
 async def build_connection_string(
     ops_test: OpsTest,
     application_name: str,
@@ -178,14 +196,24 @@ async def build_connection_string(
     )
     host = endpoints.split(",")[0].split(":")[0]
 
-    # Translate the pod hostname to an IP address.
-    model = ops_test.model.info
-    client = AsyncClient(namespace=model.name)
-    pod = await client.get(Pod, name=host.split(".")[0])
-    ip = pod.status.podIP
+    use_node_port = await is_external_connectivity_set(
+        ops_test,
+        application_name,
+        relation_name,
+        relation_id,
+    )
+
+    if use_node_port:
+        return f"dbname='{database}' user='{username}' host='{host}' port={endpoints.split(',')[0].split(':')[1]} password='{password}' connect_timeout=10"
+    else:
+        # Translate the pod hostname to an IP address.
+        model = ops_test.model.info
+        client = AsyncClient(namespace=model.name)
+        pod = await client.get(Pod, name=host.split(".")[0])
+        ip = pod.status.podIP
 
     # Build the complete connection string to connect to the database.
-    return f"dbname='{database}' user='{username}' host='{ip}' password='{password}' connect_timeout=10"
+    return f"dbname='{database}' user='{username}' host='{ip}' password='{password}' connect_timeout=10 port=6432"
 
 
 async def check_new_relation(
diff --git a/tests/integration/relations/pgbouncer_provider/test_pgbouncer_provider.py b/tests/integration/relations/pgbouncer_provider/test_pgbouncer_provider.py
index 2116f3406..6ab43a5e4 100644
--- a/tests/integration/relations/pgbouncer_provider/test_pgbouncer_provider.py
+++ b/tests/integration/relations/pgbouncer_provider/test_pgbouncer_provider.py
@@ -537,7 +537,6 @@ async def test_connection_is_possible_after_pod_deletion(ops_test: OpsTest) -> N
     connection_string = await build_connection_string(
         ops_test, DATA_INTEGRATOR_APP_NAME, relation_name="postgresql", database="test-database"
     )
-    connection_string += " port=6432"
     connection = None
     try:
         connection = psycopg2.connect(connection_string)