test: mid-through test for process_and_upload_tables

2 files changed, 39 insertions, 12 deletions

diff --git a/src/extract_lambda.py b/src/extract_lambda.py
index 6e94bba..a70ecdd 100644
--- a/src/extract_lambda.py
+++ b/src/extract_lambda.py
@@ -5,6 +5,7 @@ import csv
 from botocore.exceptions import ClientError
 import logging
 import json
+from datetime import datetime
 
 logger = logging.getLogger()
 logger.setLevel(logging.INFO)
@@ -16,7 +17,7 @@ class DBConnectionException(Exception):
         """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
@@ -71,9 +72,6 @@ def connect_to_database() -> Connection:
             host=host,
             port=port
         )
-    # except DatabaseError as e:
-    #     logger.error(f'Database error: {e}')
-    #     raise 
     except InterfaceError as i:
         logger.error(f'Interface error: {i}')
         raise DBConnectionException("Failed to connect to database")
@@ -110,14 +108,14 @@ def list_existing_s3_files(bucket_name='extract_bucket', client=boto3.client('s3
 
 
 
-def process_and_upload_tables(db, existing_files):
+def process_and_upload_tables(db, existing_files, client=boto3.client('s3')):
     """Creates a list of the tables from a database query and 
        then selects everything from each table in individual queries
        it then writes each table to CSV files and compares with the item 
-       in the existing_files dictionary with the same name. If it finds sny changes
+       in the existing_files dictionary with the same name. If it finds any changes
        to files, or new tables/files it uploads them to the s3 bucket
     """
-    client = boto3.client('s3')
+    
     tables = db.run("SELECT table_name FROM information_schema.tables WHERE table_schema='public' AND table_type='BASE TABLE';")
     
     for table in tables:
@@ -132,13 +130,13 @@ def process_and_upload_tables(db, existing_files):
             writer.writerow(column_names)
             writer.writerows(rows)
         
-        s3_key = f"{table_name}/latest.csv"
+        s3_key = f"{table_name}/{datetime.today().year}/{datetime.today().month}/{datetime.today().day}/{table_name}_{datetime.now().strftime('%H:%M:%S')}.csv"
         new_csv_content = open(csv_file_path, "r").read()
         
 
         if s3_key not in existing_files or existing_files[s3_key] != new_csv_content:
             try:
-                client.upload_file(csv_file_path, ingestion_bucket, s3_key)
+                client.upload_file(csv_file_path, 'extract_bucket', s3_key)
                 logger.info(f"Uploaded {s3_key} to S3.")
             except ClientError as e:
                 logger.error(f'Error uploading to S3: {e}')
\ No newline at end of file
diff --git a/tests/test_extract_lambda.py b/tests/test_extract_lambda.py
index 18c49fc..74d7e2c 100644
--- a/tests/test_extract_lambda.py
+++ b/tests/test_extract_lambda.py
@@ -3,7 +3,7 @@ import boto3
 from moto import mock_aws
 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
+from src.extract_lambda import list_existing_s3_files, connect_to_database, DBConnectionException, process_and_upload_tables
 import os 
 import logging
 
@@ -33,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()
@@ -80,4 +80,33 @@ class TestConnectToDatabase:
         caplog.set_level(logging.ERROR)
         with pytest.raises(DBConnectionException):
             connect_to_database()
-        assert 'Interface error' in caplog.text
\ No newline at end of file
+        assert 'Interface error' in caplog.text
+
+class TestProcessAndUploadTables:
+    def test_error_process_and_upload_tables(mock_conn, mock_config, s3_client, caplog, mocker):
+        logger = logging.getLogger()
+        logger.info('Testing now.')
+        caplog.set_level(logging.ERROR)
+
+        with patch("src.extract_lambda.Connection", autospec=True) as mock_conn:
+            mock_db = connect_to_database()
+            # need to add a table 
+            s3_key = 'dummy/2024/8/14/dummy_16:46:30.txt'
+            mock_existing_files = mocker.Mock(return_value={s3_key: 'This is a test file.' })
+            s3_client.create_bucket(Bucket='extract_bucket', 
+                        CreateBucketConfiguration={
+                            'LocationConstraint': 'eu-west-2'
+                        })
+            s3_client.upload_file('tests/dummy.txt', 'extract_bucket', s3_key)
+            process_and_upload_tables(mock_db, mock_existing_files, client=s3_client)
+
+            assert 'Error uploading to S3' in caplog.text
+
+#@pytest.mark.describe("Helpers")
+# @pytest.mark.it("Query processor returns correctly formatted dict")
+# def test_process_query():
+#     with patch("src.api.helpers.get_db_connection") as mock_conn:
+#         mock_conn().run.side_effect = db_data
+#         mock_conn().columns = sample_headers
+#         result = process_query("test query")
+#         assert result == sample_result
\ No newline at end of file