Use SQLAlchemy to connect to DB.

config/environment.yml

@@ -9,4 +9,5 @@ dependencies:
   - flake8
   - mypy
   - psycopg2
-  - python-dotenv
+  - python-dotenv
+  - sqlalchemy

setup.py

@@ -1,11 +1,17 @@
 import os
+import sqlalchemy.engine.url
 
 from dotenv import load_dotenv
 
 load_dotenv()
 
-db_host = "212.235.208.113"
-db_database = "staw"
-db_user = "staw_db"
-
 db_password = os.getenv("DB_PASSWORD")
+
+db_uri = sqlalchemy.engine.url.URL(
+    drivername='postgresql+psycopg2',
+    username="staw_db",
+    password=db_password,
+    host="212.235.208.113",
+    port=5432,
+    database="staw"
+)

test/database.py

@@ -1,19 +1,18 @@
 import unittest
 
-import psycopg2
+from sqlalchemy import create_engine
 
-from setup import db_database, db_host, db_password, db_user
+from setup import db_uri
 
 
 class DatabaseConnection(unittest.TestCase):
     def setUp(self):
-        self.conn = psycopg2.connect(
+        self.engine = create_engine(db_uri, echo=True)
-            host=db_host, user=db_user, password=db_password, database=db_database,
-        )
 
     def tearDown(self):
-        self.conn.close()
+        self.engine.dispose()
 
     def test_connection(self):
-        with self.conn.cursor() as cursor:
+        with self.engine.connect() as connection:
-            self.assertIsNotNone(cursor)
+            self.assertIsNotNone(connection)
+            connection.close()