2 files changed, 62 insertions, 18 deletions

diff --git a/src/extract_lambda.py b/src/extract_lambda.py
index dc70590..6e94bba 100644
--- a/src/extract_lambda.py
+++ b/src/extract_lambda.py
@@ -1,6 +1,5 @@
 from pg8000.native import Connection, DatabaseError, InterfaceError
-from dotenv import load_dotenv
-import os
+from dotenv import dotenv_values
 import boto3
 import csv
 from botocore.exceptions import ClientError
@@ -9,16 +8,15 @@ import json
 
 logger = logging.getLogger()
 logger.setLevel(logging.INFO)
-load_dotenv()
-
-
-database = os.getenv('database')
-user = os.getenv('user')
-password = os.getenv('password')
-host = os.getenv('host')
-port = os.getenv('port')
 
+class DBConnectionException(Exception):
+    """Wraps pg8000.native Error or DatabaseError."""
 
+    def __init__(self, e):
+        """Initialise with provided error message."""
+        self.message = str(e)
+        super().__init__(self.message)
+        
 def lambda_handler(event, context):
     """This lambda function connects to the Totesys database, lists the contents of the ingestion bucket,
        and converts all tables to CSV and if any of those tables do not exist in, or are different to the ones in s3, it uploads them
@@ -53,8 +51,19 @@ def lambda_handler(event, context):
         if db:
             db.close()
 
-def connect_to_database():
+def get_config(path: str = ".env") -> dict:
+    return dotenv_values(path)
+
+
+def connect_to_database() -> Connection:
     try:
+        config = get_config()
+        host = config["host"]
+        port = config["port"]
+        user = config["user"]
+        password = config["password"]
+        database = config["database"]
+
         return Connection(
             database=database,
             user=user,
@@ -62,12 +71,13 @@ def connect_to_database():
             host=host,
             port=port
         )
-    except DatabaseError as e:
-        logger.error(f'Database error: {e}')
-        raise
+    # except DatabaseError as e:
+    #     logger.error(f'Database error: {e}')
+    #     raise 
     except InterfaceError as i:
         logger.error(f'Interface error: {i}')
-        raise
+        raise DBConnectionException("Failed to connect to database")
+
 
 
 def list_existing_s3_files(bucket_name='extract_bucket', client=boto3.client('s3')):
diff --git a/tests/test_extract_lambda.py b/tests/test_extract_lambda.py
index 472e93a..18c49fc 100644
--- a/tests/test_extract_lambda.py
+++ b/tests/test_extract_lambda.py
@@ -1,10 +1,24 @@
 import pytest
 import boto3
 from moto import mock_aws
-from src.extract_lambda import list_existing_s3_files #process_and_upload_tables
+from unittest.mock import patch
+from unittest import TestCase
+from src.extract_lambda import list_existing_s3_files, connect_to_database, DBConnectionException #process_and_upload_tables
 import os 
 import logging
 
+@pytest.fixture(scope='class')
+def mock_config():
+    env_vars = {
+        "host": "abc",
+        "port": "5432",
+        "user": "def",
+        "password": "password",
+        "database": "db",
+    }
+    with patch("src.extract_lambda.get_config", return_value=env_vars) as mock_config:
+        yield mock_config
+
 
 @pytest.fixture(scope='class')
 def aws_credentials():
@@ -19,7 +33,7 @@ def s3_client(aws_credentials):
     with mock_aws():
         yield boto3.client('s3')
 
-class TestListExistings3Files():
+class TestListExistings3Files:
     def test_error_if_no_bucket(self, s3_client, caplog):
 
         logger = logging.getLogger()
@@ -46,4 +60,24 @@ class TestListExistings3Files():
     def test_retrieves_file_content(self, s3_client, caplog):
         result = list_existing_s3_files(client=s3_client)
 
-        assert list(result.values()) == ['This is a test file.'] 
\ No newline at end of file
+        assert list(result.values()) == ['This is a test file.'] 
+
+class TestConnectToDatabase:
+    def test_connect_to_database(mock_conn, mock_config):
+        with patch("src.extract_lambda.Connection", autospec=True) as mock_conn:  
+            connect_to_database()
+            mock_conn.assert_called_with(
+            host="abc", user="def", port="5432", password="password", database="db"
+            )
+
+    def test_database_error(self, mock_config):
+        with pytest.raises(DBConnectionException):
+            connect_to_database()
+
+    def test_logs_interface_error(self, caplog):
+        logger = logging.getLogger()
+        logger.info('Testing now.')
+        caplog.set_level(logging.ERROR)
+        with pytest.raises(DBConnectionException):
+            connect_to_database()
+        assert 'Interface error' in caplog.text
\ No newline at end of file