Remove s3_client from PredictitAPI attribute and instead supply as parameter in store_to_s3 method. Update lambdas and testing.

Author: billy-moore-98

lambda_fetch/lambda_function.py

@@ -26,7 +26,7 @@ def lambda_function(filename: str) -> None:
     bucket = os.getenv("S3_BUCKET")
     if not bucket:
         raise ValueError("S3_BUCKET environment variable is not set")
-    predictit.store_to_s3(data, bucket=bucket, filename=filename)
+    predictit.store_to_s3(s3_client, data, bucket=bucket, filename=filename)
     logging.info("Successfully stored data to S3")
 
 

src/api.py

@@ -14,7 +14,6 @@
 class PredictitAPI:
     def __init__(self, base_url="https://www.predictit.org/api/marketdata"):
         self.base_url = base_url
-        self.s3_client = boto3.client("s3")
 
     def poll_market_data(self, market_id: Optional[str] = None) -> Optional[dict]:
         """
@@ -42,14 +41,14 @@ def poll_market_data(self, market_id: Optional[str] = None) -> Optional[dict]:
             logging.error(f"An error occurred: {e}")
 
     def store_to_s3(
-        self, data: dict, bucket: Optional[str] = None, filename: Optional[str] = None
+        self, s3_client: boto3.client, data: dict, bucket: Optional[str] = None, filename: Optional[str] = None
     ):
         if not filename:
             timestamp = datetime.datetime.utcnow().strftime("%Y-%m-%dT%H-%M-%S")
             filename = f"market_data_{timestamp}.json"
         key = f"predictit/stage/{filename}"
         try:
-            self.s3_client.put_object(
+            s3_client.put_object(
                 Bucket=bucket,
                 Key=key,
                 Body=json.dumps(data),

tests/test_api.py

@@ -41,28 +41,23 @@ def test_poll_market_http_error(mock_get, caplog):
     assert "HTTP error occurred" in caplog.text
 
 
-@patch("src.api.boto3.client")
-def test_store_to_s3_success(mock_boto_client):
+def test_store_to_s3_success():
     mock_s3 = MagicMock()
-    mock_boto_client.return_value = mock_s3
     test_bucket = "test-bucket"
     test_filename = "test_file.json"
     test_data = {"markets": []}
     predictit = PredictitAPI()
     expected_key = f"predictit/stage/{test_filename}"
-    predictit.store_to_s3(test_data, bucket=test_bucket, filename=test_filename)
+    predictit.store_to_s3(mock_s3, test_data, bucket=test_bucket, filename=test_filename)
     mock_s3.put_object.assert_called_once_with(
         Bucket=test_bucket,
         Key=expected_key,
         Body='{"markets": []}',
         ContentType="application/json",
     )
 
-
-@patch("src.api.boto3.client")
-def test_store_to_s3_failure(mock_boto_client):
+def test_store_to_s3_failure():
     mock_s3 = MagicMock()
-    mock_boto_client.return_value = mock_s3
     mock_s3.put_object.side_effect = ClientError(
         error_response={"Error": {"Code": "AccessDenied", "Message": "Access Denied"}},
         operation_name="PutObject",
@@ -72,7 +67,7 @@ def test_store_to_s3_failure(mock_boto_client):
     test_filename = "test_file.json"
     predictit = PredictitAPI()
     with pytest.raises(ClientError):
-        predictit.store_to_s3(test_data, bucket=test_bucket, filename=test_filename)
+        predictit.store_to_s3(mock_s3, test_data, bucket=test_bucket, filename=test_filename)
 
     mock_s3.put_object.assert_called_once_with(
         Bucket=test_bucket,

tests/test_lambda_fetch.py

@@ -2,7 +2,7 @@
 
 from lambda_fetch.lambda_function import lambda_function, lambda_handler
 from src.api import PredictitAPI
-from unittest.mock import patch
+from unittest.mock import patch, MagicMock
 
 
 @patch("lambda_fetch.lambda_function.lambda_function")
@@ -24,17 +24,17 @@ def test_lambda_handler_no_filename(mock_lambda_func):
 
 
 @patch("lambda_fetch.lambda_function.os.getenv")
-@patch.object(PredictitAPI, "store_to_s3")
-@patch.object(PredictitAPI, "poll_market_data")
-def test_lambda_function_success(mock_poll, mock_store, mock_getenv):
-    mock_poll.return_value = {"markets": []}
+@patch("lambda_fetch.lambda_function.predictit")
+@patch("lambda_fetch.lambda_function.s3_client")
+def test_lambda_function_success(mock_s3_client, mock_predictit, mock_getenv):
     mock_getenv.return_value = "fake-bucket"
+    mock_predictit.poll_market_data.return_value = {"markets": []}
 
     lambda_function("test.json")
 
-    mock_poll.assert_called_once()
-    mock_store.assert_called_once_with(
-        {"markets": []}, bucket="fake-bucket", filename="test.json"
+    mock_predictit.poll_market_data.assert_called_once()
+    mock_predictit.store_to_s3.assert_called_once_with(
+        mock_s3_client, {"markets": []}, bucket="fake-bucket", filename="test.json"
     )