Merge pull request #102 from ajschofield/feature/load-lambda-tests

pr: feature/load lambda tests

4 files changed, 206 insertions, 45 deletions

diff --git a/src/dataframes.py b/src/dataframes.py
index f122368..2a46bd6 100644
--- a/src/dataframes.py
+++ b/src/dataframes.py
@@ -20,6 +20,28 @@ import requests
 def create_fact_sales_order(dict_of_df):
     df_sales = dict_of_df["sales_order"]
     df_sales.index.name = "sales_record_id"
+
+    df_sales["created_date"] = df_sales["created_at"].astype("datetime64[ns]").dt.date
+    df_sales["created_time"] = (
+        df_sales["created_at"].astype("datetime64[ns]").dt.floor("s").dt.time
+    )
+    df_sales["last_updated_date"] = (
+        df_sales["last_updated"].astype("datetime64[ns]").dt.date
+    )
+    df_sales["last_updated_time"] = (
+        df_sales["last_updated"].astype("datetime64[ns]").dt.floor("s").dt.time
+    )
+    df_sales["agreed_delivery_date"] = pd.to_datetime(
+        df_sales["agreed_delivery_date"], format="%Y-%m-%d"
+    )
+    df_sales["agreed_payment_date"] = pd.to_datetime(
+        df_sales["agreed_payment_date"], format="%Y-%m-%d"
+    )
+    df_sales = df_sales.drop(labels=["created_at", "last_updated"], axis=1)
+
+    df_sales.reset_index(inplace=True)
+    return df_sales
+
     df_sales["created_date"] = df_sales["created_at"].astype("datetime64[ns]").dt.date
     df_sales["created_time"] = (
         df_sales["created_at"].astype("datetime64[ns]").dt.floor("s").dt.time
@@ -88,6 +110,7 @@ def create_fact_payment(dict_of_df):
         df_payment["payment_date"], format="%Y-%m-%d"
     )
     df_payment = df_payment.drop(labels=["created_at", "last_updated"], axis=1)
+
     df_payment.reset_index(inplace=True)
     return df_payment
 
@@ -115,18 +138,24 @@ def create_dim_location(dict_of_df):
 
 
 def create_dim_counterparty(dict_of_df):
-    df_prefixed_address = dict_of_df["address"].add_prefix(
-        "counterparty_legal_", axis=1
+    df_prefixed_address = (
+        dict_of_df["address"]
+        .drop(labels=["created_at", "last_updated"], axis=1)
+        .add_prefix("counterparty_legal_", axis=1)
     )
     df_cp = pd.merge(
         dict_of_df["counterparty"],
         df_prefixed_address,
         left_on="legal_address_id",
         right_on="counterparty_legal_address_id",
-        how="outer",
+        how="inner",
     )
     df_cp.drop(
-        columns=["legal_address_id", "counterparty_legal_address_id"], inplace=True
+        columns=[
+            "legal_address_id",
+            "counterparty_legal_address_id",
+        ],
+        inplace=True,
     )
     return df_cp
 
diff --git a/src/load_lambda.py b/src/load_lambda.py
index 6e6bc80..7339ab9 100644
--- a/src/load_lambda.py
+++ b/src/load_lambda.py
@@ -5,7 +5,7 @@ import pyarrow.parquet as pq
 from io import BytesIO
 import logging
 import json
-from src.extract_lambda import retrieve_secrets
+import traceback
 from sqlalchemy import create_engine
 
 
@@ -24,33 +24,39 @@ logging.getLogger("botocore").setLevel(logging.INFO)
 def lambda_handler(event, context):
     try:
         uploaded_tables = upload_dfs_to_database()
-        if not uploaded_tables["uploaded"]:
+        if uploaded_tables["not_uploaded"]:
             return {
                 "statusCode": 200,
                 "body": json.dumps("No dataframes were uploaded."),
             }
-        return {
-            "statusCode": 200,
-            "body": json.dumps(
-                f"""The following dataframes were uploaded successfully: 
-                {uploaded_tables["uploaded"]} ."""
-            ),
-        }
+        elif uploaded_tables["uploaded"]:
+            return {
+                "statusCode": 200,
+                "body": json.dumps(
+                    f"""The following dataframes were uploaded successfully: 
+                    {uploaded_tables["uploaded"]} ."""
+                ),
+            }
+        else:
+            logger.error(f"error")
+            return {"error"}
     except Exception as e:
-        logger.error(f"Error: {e}", exc_info=True)
-        return {"statusCode": 500, "body": json.dumps("Internal server error.")}
+        logger.error({e})
+        return {"statusCode": 500, "body": {e}}
 
 
-def retrieve_secrets():
-    secret_name = "bentley-RDS-credentials"
+def retrieve_secrets(client=None, secret_name=None):
+    session = boto3.session.Session()
     region_name = "eu-west-2"
 
-    # Create a Secrets Manager client
-    session = boto3.session.Session()
-    client = session.client(service_name="secretsmanager", region_name=region_name)
+    if secret_name == None:
+        secret_name = "bentley-RDS-credentials"
+    if client == None:
+        client = session.client(service_name="secretsmanager", region_name=region_name)
 
     try:
         get_secret_value_response = client.get_secret_value(SecretId=secret_name)
+        print(get_secret_value_response)
     except ClientError as e:
         logger.error(f"Failed to retrieve secret {secret_name}: {str(e)}")
         raise e
@@ -64,9 +70,12 @@ def retrieve_secrets():
 # connect to database, slightly different way of doing it, to allow manipulation through pandas
 
 
-def connect_to_db_and_return_engine():
+def connect_to_db_and_return_engine(sm_secret=None):
+    if sm_secret is None:
+        sm_secret = json.loads(retrieve_secrets())
+
     try:
-        secrets = json.loads(retrieve_secrets())
+        secrets = sm_secret
         host = secrets["host"]
         port = secrets["port"]
         user = secrets["user"]
@@ -162,14 +171,15 @@ def upload_dfs_to_database():
     ]
 
     for file_name, df in dict_of_dfs.items():
+        print(df)
         if file_name in immutable_df_dict:
             table_name = file_name.split(".")[0]
+            print(table_name, "<<<<<")
             try:
                 df.to_sql(
                     table_name,
                     con=db_engine,
-                    schema="project_team_2",
-                    if_exists="overwrite",
+                    if_exists="append",
                     index=False,
                 )
                 upload_status["uploaded"].append(table_name)
@@ -183,7 +193,7 @@ def upload_dfs_to_database():
                     table_name,
                     con=db_engine,
                     schema="project_team_2",
-                    if_exists="overwrite",
+                    if_exists="append",
                     index=False,
                 )
                 upload_status["uploaded"].append(table_name)
@@ -195,3 +205,7 @@ def upload_dfs_to_database():
             logger.error(f"{file_name} does not correspond with table in database")
     db_engine.dispose()
     return upload_status
+
+
+if __name__ == "__main__":
+    lambda_handler(None, None)
diff --git a/tests/test_dataframes.py b/tests/test_dataframes.py
index c9ff43f..ea7bad1 100644
--- a/tests/test_dataframes.py
+++ b/tests/test_dataframes.py
@@ -54,7 +54,8 @@ class TestCreateDimStaff:
             "email_address": ["Hello", "Bye"],
             "department_id": ["Hello", "Bye"],
         }
-        test_df = {"staff": pd.DataFrame(data=d), "department": pd.DataFrame(data=d2)}
+        test_df = {"staff": pd.DataFrame(
+            data=d), "department": pd.DataFrame(data=d2)}
         result = create_dim_staff(test_df)
         assert isinstance(result, pd.DataFrame)
 
@@ -71,7 +72,10 @@ class TestCreateDimStaff:
             "email_address": ["Hello", "Bye"],
             "department_id": ["Hello", "Bye"],
         }
-        test_df = {"staff": pd.DataFrame(data=d), "department": pd.DataFrame(data=d2)}
+
+        test_df = {"staff": pd.DataFrame(
+            data=d), "department": pd.DataFrame(data=d2)}
+
         result = create_dim_staff(test_df)
         expected_d = {
             "staff_id": ["Hello", "Bye"],
@@ -88,7 +92,9 @@ class TestCreateDimStaff:
 
 class TestCreatePaymentType:
     def test_create_dim_payment_type_returns_correct_columns_and_values(self):
-        d = {"payment_type_id": ["Hello", "Bye"], "payment_type_name": ["Hello", "Bye"]}
+        d = {"payment_type_id": ["Hello", "Bye"],
+             "payment_type_name": ["Hello", "Bye"]}
+
         test_df = {"payment_type": pd.DataFrame(data=d)}
         result = create_dim_payment_type(test_df)
         expected_columns = ["payment_type_id", "payment_type_name"]
@@ -180,11 +186,14 @@ class TestCreateDimDate:
             index=[0],
         )
         df_two = pd.DataFrame(
-            data={"updated_date": dt(2020, 5, 17), "created_date": dt(2021, 9, 13)},
+            data={"updated_date": dt(2020, 5, 17),
+                  "created_date": dt(2021, 9, 13)},
             index=[0],
         )
         df_three = pd.DataFrame(
-            data={"updated_date": dt(2022, 5, 17), "created_date": dt(2023, 5, 13)},
+            data={"updated_date": dt(2022, 5, 17),
+                  "created_date": dt(2023, 5, 13)},
+
             index=[0],
         )
         expected_df = pd.DataFrame(
@@ -214,7 +223,8 @@ class TestCreateDimDate:
                     mock_fso.return_value = df_three
                     result = create_dim_date({"dum": 0})
                     result.reset_index(inplace=True, drop=True)
-                    assert result.eq(expected_df, axis="columns").all(axis=None)
+                    assert result.eq(
+                        expected_df, axis="columns").all(axis=None)
 
 
 class TestCreateDimLocation:
@@ -222,7 +232,9 @@ class TestCreateDimLocation:
         dict_df = {
             "address": pd.DataFrame(
                 data=[["some_time", "some_other_time", 1, "SE18 9QO"]],
-                columns=["created_at", "last_updated", "address_id", "postal_code"],
+                columns=["created_at", "last_updated",
+                         "address_id", "postal_code"],
+
             )
         }
         result = create_dim_location(dict_df)
@@ -287,5 +299,7 @@ class TestCreateFactPayment:
         for col in list(result.columns):
             assert col in expected_cols
         for col in expected_cols:
-            if "_date" or "_time" in col:
-                assert result[col].dtype == "O"
+
+
+if "_date" or "_time" in col:
+    assert result[col].dtype == "O"
diff --git a/tests/test_load_lambda.py b/tests/test_load_lambda.py
index 88c71e4..65106f7 100644
--- a/tests/test_load_lambda.py
+++ b/tests/test_load_lambda.py
@@ -3,22 +3,18 @@ import pyarrow.parquet as pq
 from io import BytesIO
 from moto import mock_aws
 import boto3
+import botocore.exceptions
 import os
 import pytest
-from src.load_lambda import (
-    lambda_handler,
-    connect_to_db_and_return_engine,
-    get_transform_bucket,
-    convert_parquet_files_to_dfs,
-    upload_dfs_to_database,
-)
+from src.load_lambda import *
+import tempfile
 
 
 @pytest.fixture(scope="class")
 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"
 
@@ -29,16 +25,89 @@ def mock_s3_client(aws_credentials):
         yield boto3.client("s3")
 
 
+@pytest.fixture(scope="class")
+def mock_sm_client(aws_credentials):
+    with mock_aws():
+        yield boto3.client("secretsmanager")
+
+
 class TestLambdaHandler:
-    pass
+    def test_lambda_handler_returns_200_and_table_name_if_uploaded(self, mocker):
+        mocker.patch(
+            "src.load_lambda.upload_dfs_to_database",
+            return_value={"uploaded": ["table_one", "table_two"], "not_uploaded": []},
+        )
+        result = lambda_handler(None, None)
+        assert result["statusCode"] == 200
+        assert "table_one" in result["body"]
+        assert "table_two" in result["body"]
+
+    def test_lambda_handler_returns_200_and_table_name_if_not_uploaded(self, mocker):
+        mocker.patch(
+            "src.load_lambda.upload_dfs_to_database",
+            return_value={"uploaded": [], "not_uploaded": ["table_one"]},
+        )
+        result = lambda_handler(None, None)
+        assert result["statusCode"] == 200
+        assert "No dataframes were uploaded" in result["body"]
+
+    def test_lambda_handler_returns_error_if_both_lists_empty(self, mocker):
+        mocker.patch(
+            "src.load_lambda.upload_dfs_to_database",
+            return_value={"uploaded": [], "not_uploaded": []},
+        )
+
+        result = lambda_handler(None, None)
+
+        assert result == {"error"}
 
 
 class TestRetrieveSecrets:
-    pass
+    def test_retrieve_secrets_returns_dictionary(self, mock_sm_client):
+        secret = {
+            "cohort_id": "test_cohort_id",
+            "user": "test_user_id",
+            "password": "test_password",
+            "host": "test_host",
+            "database": "test_database",
+            "port": "test_port",
+        }
+
+        secret_name = "test_secret"
+
+        mock_sm_client.create_secret(Name=secret_name, SecretString=json.dumps(secret))
+
+        result = json.loads(retrieve_secrets(mock_sm_client, secret_name))
+
+        assert isinstance(result, dict)
+
+    def test_retrieve_secrets_returns_correct_keys_and_values(self, mock_sm_client):
+        secret_name = "test_secret"
+
+        result = json.loads(retrieve_secrets(mock_sm_client, secret_name))
+
+        assert result["user"] == "test_user_id"
+        assert result["password"] == "test_password"
+
+    def test_retrieve_secrets_returns_client_error_if_no_secret(self, mock_sm_client):
+        secret_name = "another_test_secret"
+
+        with pytest.raises(botocore.exceptions.ClientError) as error:
+            retrieve_secrets(mock_sm_client, secret_name)
 
 
 class TestConnectToDBAndReturnEngine:
-    pass
+    def test_returns_unsuccessful_connection_when_wrong_credentials(self):
+        sm_secret = {
+            "host": "host",
+            "port": "port",
+            "user": "user",
+            "password": "password",
+            "database": "database",
+        }
+
+        with pytest.raises(Exception):
+            connect_to_db_and_return_engine(json.dumps(sm_secret))
 
 
 class TestGetTransformBucket:
@@ -87,6 +156,41 @@ class TestConvertParquetToDfs:
     #     result = convert_parquet_files_to_dfs(bucket_name="transform_bucket", client=mock_s3_client)
     #     assert "dim_staff" in result
 
+    def test_function_returns_dictionary_with_file_key_and_dataframe(
+        self, mock_s3_client
+    ):
+        with tempfile.TemporaryDirectory() as tmp:
+            d = {
+                "test": ["Hello", "Bye"],
+                "design_id": ["Hello", "Bye"],
+                "design_name": ["Hello", "Bye"],
+                "file_name": ["Hello", "Bye"],
+                "file_location": ["Hello", "Bye"],
+                "Hello": ["Hello", "Bye"],
+            }
+
+            test_df = pd.DataFrame(data=d)
+
+            path = os.path.join(tmp, "test_parquet.parquet")
+
+            test_df.to_parquet(path, engine="pyarrow")
+
+            with open(path, "rb") as p:
+                mock_s3_client.put_object(
+                    Bucket="transform_bucket", Key="test_parquet.parquet", Body=p.read()
+                )
+
+            result = convert_parquet_files_to_dfs(
+                bucket_name="transform_bucket", client=mock_s3_client
+            )
+
+            assert "test_parquet.parquet" in result
+
+            pd.testing.assert_frame_equal(result["test_parquet.parquet"], test_df)
+
 
 class TestUploadDfsToDatabase:
+    # Full success test
+    # Partial success test
+    # Failure test
     pass