wip: update TestLambdaHandler & lambda_handler function

2 files changed, 20 insertions, 11 deletions

diff --git a/src/load_lambda.py b/src/load_lambda.py
index 11d1d70..39fa27d 100644
--- a/src/load_lambda.py
+++ b/src/load_lambda.py
@@ -23,18 +23,21 @@ 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"]} ."""
-            ),
-        }
+
+        if uploaded_tables["uploaded"]:
+            return {
+                "statusCode": 200,
+                "body": json.dumps(
+                    f"""The following dataframes were uploaded successfully: 
+                    {uploaded_tables["uploaded"]} ."""
+                ),
+            }
+
     except Exception as e:
         logger.error(f"Error: {e}", exc_info=True)
         return {"statusCode": 500, "body": json.dumps("Internal server error.")}
diff --git a/tests/test_load_lambda.py b/tests/test_load_lambda.py
index a29b75a..9286e48 100644
--- a/tests/test_load_lambda.py
+++ b/tests/test_load_lambda.py
@@ -35,7 +35,7 @@ class TestLambdaHandler:
     def test_lambda_handler_returns_success(self, mocker):
         mocker.patch(
             "src.load_lambda.upload_dfs_to_database",
-            return_value={"uploaded": ["table_one", "table_two"]},
+            return_value={"uploaded": ["table_one", "table_two"], "not_uploaded": []},
         )
         result = lambda_handler(None, None)
         assert result["statusCode"] == 200
@@ -45,14 +45,20 @@ class TestLambdaHandler:
     def test_lambda_handler_does_not_upload_anything(self, mocker):
         mocker.patch(
             "src.load_lambda.upload_dfs_to_database",
-            return_value={"uploaded": []},
+            return_value={"uploaded": [], "not_uploaded": []},
         )
         result = lambda_handler(None, None)
         assert result["statusCode"] == 200
         assert "No dataframes were uploaded" in result["body"]
 
     def test_lambda_handler_returns_exception(self, mocker):
-        pass
+        mocker.patch(
+            "src.load_lambda.upload_dfs_to_database",
+            return_value={"test": []},
+        )
+
+        with pytest.raises(Exception):
+            lambda_handler(None, None)
 
 
 class TestRetrieveSecrets: