change(test_postgresql): clean, comment, format

tests/test_postgresql.py

@@ -12,7 +12,10 @@ import subprocess
 import sys
 import time
 
-import psycopg2
+try:
+    import psycopg2
+except ImportError:
+    sys.exit(2)
 
 GREEN = "\033[92m"
 YELLOW = "\033[93m"
@@ -134,20 +137,6 @@ def get_nodes(conf: dict) -> dict:
     return servers
 
 
-def check_odd_number(number: int) -> bool:
-    """Check if we have an odd number of nodes, ensuring we can have a quorum.
-
-    :param number: The number of nodes in the cluster
-    :type number: int
-    :return: Wether it as an odd number or not
-    :rtype: bool
-    """
-
-    modulo = number % 2
-
-    return modulo != 0
-
-
 def get_node_state(host: str, port: int) -> str:
     """Get the curent state of node from its RepHACheck daemon.
 
@@ -170,87 +159,55 @@ def get_node_state(host: str, port: int) -> str:
     return state
 
 
-def check_primary(nodes: dict) -> tuple:
-    """Check if we have a primary in the nodes.
-
-    :param nodes: The dictionary containing nodes and their informations
-    :type nodes: dict
-    :return: Wether the nodes list contains a primary server
-    :rtype: tuple
-    """
-
-    for node in nodes.keys():
-        host = nodes[node]["host"]
-        port = int(nodes[node]["rephacheck"])
-        if get_node_state(host, port) == "primary":
-            return True, node
-
-    return False, None
-
-
-def check_standby(nodes: dict) -> tuple:
-    """Check if we have a standby in the nodes.
-
-    :param nodes: The dictionary containing nodes and their informations
-    :type nodes: dict
-    :return: Wether the nodes list contains a standby server
-    :rtype: tuple
-    """
-
-    for node in nodes.keys():
-        host = nodes[node]["host"]
-        port = int(nodes[node]["rephacheck"])
-        if get_node_state(host, port) == "standby":
-            return True, node
-
-    return False, None
-
-
-def check_witness(nodes: dict) -> tuple:
-    """Check if we have a witness in the nodes.
+def check_fenced(nodes: dict) -> tuple:
+    """Check if the cluster have a fenced node.
 
     :param nodes: The dictionary containing nodes and their informations
     :type nodes: dict
-    :return: Wether the nodes list contains a witness server
+    :return: Wether the nodes list contains a fenced server
     :rtype: tuple
     """
 
     for node in nodes.keys():
         host = nodes[node]["host"]
         port = int(nodes[node]["rephacheck"])
-        if get_node_state(host, port) == "witness":
+        if get_node_state(host, port) == "fenced":
             return True, node
 
     return False, None
 
 
-def check_fenced(nodes: dict) -> tuple:
-    """Check if the cluster have a fenced node.
+def check_listen(host: str, port: int) -> bool:
+    """Check if server is listening (TCP only).
 
-    :param nodes: The dictionary containing nodes and their informations
-    :type nodes: dict
-    :return: Wether the nodes list contains a fenced server
-    :rtype: tuple
+    :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
     """
 
-    for node in nodes.keys():
-        host = nodes[node]["host"]
-        port = int(nodes[node]["rephacheck"])
-        if get_node_state(host, port) == "fenced":
-            return True, node
+    # 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 False, None
+    return result == 0
 
 
 def check_psql(db_conn: dict, query: str) -> bool:
     """Check if we can write data on this node.
 
-    :param db_conn: Connection details
-    :type db_conn: str
+    :param db_conn: Database connection parameters
+    :type db_conn: dict
+    :param query: Query to execute
+    :type query: str
     :return: Wether the query can be executed or not
     :rtype: bool
     """
 
+    # build database connection uri
     if "password" in db_conn:
         uri = "postgresql://{}:{}@{}:{}/{}".format(
             db_conn["user"],
@@ -262,8 +219,10 @@ def check_psql(db_conn: dict, query: str) -> bool:
     else:
         uri = "postgresql:///{}".format(db_conn["dbname"])
 
+    # format command
     command = ["su -l postgres -c \"psql {} -c '{}'\"".format(uri, query)]
 
+    # execute
     try:
         subprocess.check_output(command, shell=True)
     except subprocess.CalledProcessError:
@@ -275,9 +234,9 @@ def check_psql(db_conn: dict, query: str) -> bool:
 def check_replication(primary: dict, standby: dict) -> bool:
     """Check if replication is working between the primary and standby servers.
 
-    :param primary: Connection details for primary server
+    :param primary: Database connection parameters for primary server
     :type primary: dict
-    :param standby: Connection details for standby server
+    :param standby: Database connection parameters for standby server
     :type standby: dict
     :return: Wether replication between primary/stanbdy is working or not
     :rtype: bool
@@ -311,25 +270,6 @@ def check_replication(primary: dict, standby: dict) -> bool:
     return True
 
 
-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_ha(db_conn: dict, errors: int = 0) -> int:
     """Run all tests for a highly-available setup.
 
@@ -402,8 +342,6 @@ def check_local(db_conn: dict, errors: int = 0) -> int:
     :type db_conn: dict
     :param errors: Error counter, defaults to 0
     :param errors: int, optional
-    :param warnings: Warning counter, defaults to 0
-    :param warnings: int, optional
     :return: Number of errors
     :rtype: int
     """
@@ -422,8 +360,8 @@ def check_local(db_conn: dict, errors: int = 0) -> int:
 
     # check read
     print("Checking read operation:")
-    rquery = "SELECT 1;"
-    if not check_psql(db_conn, rquery):
+    read_query = "SELECT 1;"
+    if not check_psql(db_conn, read_query):
         error("Cannot read data on {}@{}:{}".format(db_user, db_host, db_port))
         errors += 1
     else:
@@ -431,8 +369,8 @@ def check_local(db_conn: dict, errors: int = 0) -> int:
 
     # check write
     print("Checking write operation:")
-    wquery = "CREATE TABLE es_test (id serial PRIMARY KEY);"
-    if not check_psql(db_conn, wquery):
+    write_query = "CREATE TABLE es_test (id serial PRIMARY KEY);"
+    if not check_psql(db_conn, write_query):
         error("Cannot write data on {}@{}:{}".format(db_user, db_host, db_port))
         errors += 1
     else:
@@ -470,14 +408,12 @@ def main():
     if db_pass:
         db_conn.update({"password": db_pass})
 
-    # determine if HA setup
+    # determine if HA setup and run according tests
     if is_ha(db_port):
         print("This setup is using a HA database")
-
         errors = check_ha(db_conn)
     else:
         print("This setup is using a local database")
-
         errors = check_local(db_conn)
 
     if errors: