Fix postgresql test in HA environment

tests/test_postgresql.py

@@ -23,18 +23,31 @@ sys.path.append(str(Path(__file__).parents[1].resolve()))
 import utils as u  # noqa: E402
 from utils_lib.apt import Apt  # noqa: E402
 
+def check_listen(host: str, port: int) -> bool:
+    """Check if server is listening (TCP only).
 
-def is_ha(port: int) -> bool:
-    """Check wether this setup is using higlhy-available databases.
-
-    :param port: Port number
+    :param host: The hostname or IP address to bind
+    :param port: The port number to bind
+    :type host: str
     :type port: int
-    :return: Wether it is a highly-available setup or not
+    :return: Wether the `host` is listening on TCP/`port`
     :rtype: bool
     """
 
-    return port == 54321
+    # try to connect to the port used by psql-primary frontend
+    client = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+    result = client.connect_ex((host, port))
+    client.close()
+
+    return result == 0
+
+def is_ha() -> bool:
+    """Check wether this setup is using higlhy-available databases.
 
+    :return: Wether it is a highly-available setup or not
+    :rtype: bool
+    """
+    return check_listen("127.0.0.1", 54321)
 
 def get_haproxy_conf(path: str = "/etc/haproxy/haproxy.cfg") -> dict:
     """Get HAProxy configuration in a dictionary.
@@ -148,25 +161,6 @@ def check_fenced(nodes: dict) -> tuple:
     return False, None
 
 
-def check_listen(host: str, port: int) -> bool:
-    """Check if server is listening (TCP only).
-
-    :param host: The hostname or IP address to bind
-    :param port: The port number to bind
-    :type host: str
-    :type port: int
-    :return: Wether the `host` is listening on TCP/`port`
-    :rtype: bool
-    """
-
-    # try to connect to the port used by psql-primary frontend
-    client = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-    result = client.connect_ex((host, port))
-    client.close()
-
-    return result == 0
-
-
 def check_psql(db_conn: dict, query: str) -> tuple:
     """Check if we can write data on this node.
 
@@ -362,18 +356,27 @@ def check_local(db_conn: dict, errors: int = 0) -> int:
     else:
         u.success("can read from {}@{}:{}".format(user, host, port))
 
+    # get replication state if available
+    if check_listen("127.0.0.1", 8543):
+        state = get_node_state("127.0.0.1", 8543)
+    else:
+        state = 'primary'
+
     # check write
     print("Checking write operation:")
-    rand = uuid.uuid4().hex
-    write_query = "CREATE TABLE es_test_{} (id serial PRIMARY KEY);".format(rand)
-    status, info = check_psql(db_conn, write_query)
-    if not status:
-        u.error("cannot write on {}@{}:{} ({})".format(user, host, port, info))
-        errors += 1
+    if state != "primary":
+        u.info("this database is in {} state".format(state))
     else:
-        u.success("can write on {}@{}:{}".format(user, host, port))
-        # remove test table
-        check_psql(db_conn, "DROP TABLE es_test_{};".format(rand))
+        rand = uuid.uuid4().hex
+        write_query = "CREATE TABLE es_test_{} (id serial PRIMARY KEY);".format(rand)
+        status, info = check_psql(db_conn, write_query)
+        if not status:
+            u.error("cannot write on {}@{}:{} ({})".format(user, host, port, info))
+            errors += 1
+        else:
+            u.success("can write on {}@{}:{}".format(user, host, port))
+            # remove test table
+            check_psql(db_conn, "DROP TABLE es_test_{};".format(rand))
 
     return errors
 
@@ -390,7 +393,7 @@ def main():
 
     # get database configuration
     db_host = conf.get("DB_HOST") if conf.get("DB_HOST") else "127.0.0.1"
-    db_port = int(conf.get("DB_PORT")) if conf.get("DB_PORT") else 5432
+    db_port = 54321 if is_ha() else 5432
     db_user = conf.get("DB_USER") if conf.get("DB_USER") else "postgres"
     db_pass = conf.get("DB_PG_ROOT_PWD")
     db_conn = {"dbname": db_user, "host": db_host, "port": db_port, "user": db_user}
@@ -399,7 +402,7 @@ def main():
 
     # determine if HA setup and run according tests
     print("Checking availibility mode:")
-    if is_ha(db_port):
+    if is_ha():
         u.info("this setup is using a HA database")
         errors = check_ha(db_conn)
     else: