Merge pull request #64 from ajschofield/extract-lambda-code-update

pr: extract lambda code update

2 files changed, 22 insertions, 45 deletions

diff --git a/src/extract_lambda.py b/src/extract_lambda.py
index 4168e27..217efdb 100644
--- a/src/extract_lambda.py
+++ b/src/extract_lambda.py
@@ -90,6 +90,7 @@ def extract_bucket(client=boto3.client("s3")):
     extract_bucket_filter = [
         bucket["Name"] for bucket in response["Buckets"] if "extract" in bucket["Name"]
     ]
+
     return extract_bucket_filter[0]
 
 
diff --git a/tests/test_extract_lambda.py b/tests/test_extract_lambda.py
index 7707cbf..a43ae0a 100644
--- a/tests/test_extract_lambda.py
+++ b/tests/test_extract_lambda.py
@@ -1,3 +1,5 @@
+import boto3.exceptions
+import botocore.exceptions
 import pytest
 import boto3
 from moto import mock_aws
@@ -24,7 +26,9 @@ def mock_config():
         "password": "password",
         "database": "db",
     }
-    with patch("src.extract_lambda.get_config", return_value=env_vars) as mock_config:
+    with patch(
+        "src.extract_lambda.retrieve_secrets", return_value=env_vars
+    ) as mock_config:
         yield mock_config
 
 
@@ -32,7 +36,7 @@ def mock_config():
 def aws_credentials():
     os.environ["AWS_ACCESS_KEY_ID"] = "testing"
     os.environ["AWS_SECRET_ACCESS_KEY"] = "testing"
-    os.environ["AWS_SECURIT_TOKEN"] = "testing"
+    os.environ["AWS_SECURITY_TOKEN"] = "testing"
     os.environ["AWS_SESSION_TOKEN"] = "testing"
     os.environ["AWS_DEFAULT_REGION"] = "eu-west-2"
 
@@ -43,6 +47,15 @@ def s3_client(aws_credentials):
         yield boto3.client("s3")
 
 
+@pytest.fixture(scope="class")
+def s3_mock_bucket(s3_client):
+    bucket = s3_client.create_bucket(
+        Bucket="extract_bucket",
+        CreateBucketConfiguration={"LocationConstraint": "eu-west-2"},
+    )
+    return bucket
+
+
 class TestLambdaHandler:
     def test_lambda_handler_files_processed_and_uploaded_successfully(self, mocker):
         mock_db = MagicMock()
@@ -135,27 +148,18 @@ class TestListExistingS3Files:
         list_existing_s3_files(client=s3_client)
         assert "Error listing S3 objects" in caplog.text
 
-    def test_error_if_bucket_is_empty(self, s3_client, caplog):
-        s3_client.create_bucket(
-            Bucket="extract_bucket",
-            CreateBucketConfiguration={"LocationConstraint": "eu-west-2"},
-        )
-        list_existing_s3_files(client=s3_client)
+    def test_error_if_bucket_is_empty(self, s3_client, caplog, s3_mock_bucket):
+        list_existing_s3_files("extract_bucket", client=s3_client)
         assert "The bucket is empty" in caplog.text
 
-    def test_error_retrieving_object(self, s3_client, caplog):
+    def test_retrieves_file_content(self, s3_client, caplog, s3_mock_bucket):
         s3_client.upload_file("tests/dummy.txt", "extract_bucket", "dummy.txt")
-        list_existing_s3_files(bucket_name="test_bucket", client=s3_client)
-
-        assert "Error retrieving S3 object " in caplog.text
-
-    def test_retrieves_file_content(self, s3_client, caplog):
-        result = list_existing_s3_files(client=s3_client)
-
+        result = list_existing_s3_files("extract_bucket", client=s3_client)
         assert list(result.values()) == ["This is a test file."]
 
 
 class TestConnectToDatabase:
+    # had mock_config in param
     def test_connect_to_database(mock_conn, mock_config):
         with patch("src.extract_lambda.Connection", autospec=True) as mock_conn:
             connect_to_database()
@@ -163,7 +167,7 @@ class TestConnectToDatabase:
                 host="abc", user="def", port="5432", password="password", database="db"
             )
 
-    def test_database_error(self, mock_config):
+    def test_database_error(self, mock_config):  # had mock_config in param
         with pytest.raises(DBConnectionException):
             connect_to_database()
 
@@ -174,31 +178,3 @@ class TestConnectToDatabase:
         with pytest.raises(DBConnectionException):
             connect_to_database()
         assert "Interface error" in caplog.text
-
-
-"""
-class TestProcessAndUploadTables:
-    def test_error_process_and_upload_tables(mock_conn, mock_config, s3_client, caplog):
-        logger = logging.getLogger()
-        logger.info('Testing now.')
-        caplog.set_level(logging.ERROR)
-        ####
-        queries = ["SELECT table_name FROM information_schema.tables WHERE table_schema='public' AND table_type='BASE TABLE';",
-                    "SELECT * FROM Fruits;",
-                    "SELECT column_name FROM INFORMATION_SCHEMA.COLUMNS where table_name = 'Fruits'"]
-        return_values = [[['Fruits']],
-                        [['Vegetable','Sour','Green'],['Berry','Sweet','Red']],
-                        [['Food_type'],['Flavour'],['Colour']]]
-        vals = dict(zip(queries,return_values))
-
-        ####
-        with patch('src.extract_lambda.connect_to_database') as mock_db:
-            mock_db().run.side_effects = return_values
-            s3_key = 'Fruits/2024/08/15/Fruits_16:46:30.csv'
-            existing_files = {s3_key: 'Food_type,Flavour,Colour\nFruit,Sour,Green\nBerry,Sweet,Red'}
-            s3_client.create_bucket(Bucket='extract_bucket', 
-                        CreateBucketConfiguration={'LocationConstraint': 'eu-west-2'})
-            s3_client.upload_file('tests/dummy_identical.csv', 'extract_bucket', s3_key)
-            process_and_upload_tables(mock_db(), existing_files, client=s3_client)
-            assert 'No new data.' in caplog.text
-"""