1 files changed, 87 insertions, 0 deletions
diff --git a/tests/features/db_utils.py b/tests/features/db_utils.py
new file mode 100644
index 0000000..595c6c2
--- /dev/null
+++ b/tests/features/db_utils.py
@@ -0,0 +1,87 @@
+from psycopg import connect
+
+
+def create_db(
+    hostname="localhost", username=None, password=None, dbname=None, port=None
+):
+    """Create test database.
+
+    :param hostname: string
+    :param username: string
+    :param password: string
+    :param dbname: string
+    :param port: int
+    :return:
+
+    """
+    cn = create_cn(hostname, password, username, "postgres", port)
+
+    cn.autocommit = True
+    with cn.cursor() as cr:
+        cr.execute(f"drop database if exists {dbname}")
+        cr.execute(f"create database {dbname}")
+
+    cn.close()
+
+    cn = create_cn(hostname, password, username, dbname, port)
+    return cn
+
+
+def create_cn(hostname, password, username, dbname, port):
+    """
+    Open connection to database.
+    :param hostname:
+    :param password:
+    :param username:
+    :param dbname: string
+    :return: psycopg2.connection
+    """
+    cn = connect(
+        host=hostname, user=username, dbname=dbname, password=password, port=port
+    )
+
+    print(f"Created connection: {cn.info.get_parameters()}.")
+    return cn
+
+
+def pgbouncer_available(hostname="localhost", password=None, username="postgres"):
+    cn = None
+    try:
+        cn = create_cn(hostname, password, username, "pgbouncer", 6432)
+        return True
+    except:
+        print("Pgbouncer is not available.")
+    finally:
+        if cn:
+            cn.close()
+    return False
+
+
+def drop_db(hostname="localhost", username=None, password=None, dbname=None, port=None):
+    """
+    Drop database.
+    :param hostname: string
+    :param username: string
+    :param password: string
+    :param dbname: string
+    """
+    cn = create_cn(hostname, password, username, "postgres", port)
+
+    # Needed for DB drop.
+    cn.autocommit = True
+
+    with cn.cursor() as cr:
+        cr.execute(f"drop database if exists {dbname}")
+
+    close_cn(cn)
+
+
+def close_cn(cn=None):
+    """
+    Close connection.
+    :param connection: psycopg2.connection
+    """
+    if cn:
+        cn_params = cn.info.get_parameters()
+        cn.close()
+        print(f"Closed connection: {cn_params}.")