Merge branch 'development' of https://github.com/ajschofield/de-project-bentley into development

1 files changed, 115 insertions, 11 deletions

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