Merge pull request #1376 from GSA/notify-api-1351

increase code coverage to 94%

.ds.baseline

@@ -209,7 +209,7 @@
         "filename": "tests/app/aws/test_s3.py",
         "hashed_secret": "67a74306b06d0c01624fe0d0249a570f4d093747",
         "is_verified": false,
-        "line_number": 29,
+        "line_number": 40,
         "is_secret": false
       }
     ],

.github/workflows/checks.yml

@@ -63,7 +63,7 @@ jobs:
         NOTIFY_E2E_TEST_PASSWORD: ${{ secrets.NOTIFY_E2E_TEST_PASSWORD }}
     - name: Check coverage threshold
       # TODO get this back up to 95
-      run: poetry run coverage report -m --fail-under=93
+      run: poetry run coverage report -m --fail-under=94
 
   validate-new-relic-config:
     runs-on: ubuntu-latest

Makefile

@@ -84,7 +84,7 @@ test: ## Run tests and create coverage report
 	poetry run coverage run --omit=*/migrations/*,*/tests/* -m pytest --maxfail=10
 
     ## TODO set this back to 95 asap
-	poetry run coverage report -m --fail-under=93
+	poetry run coverage report -m --fail-under=94
 	poetry run coverage html -d .coverage_cache
 
 .PHONY: py-lock

app/aws/s3.py

@@ -70,9 +70,13 @@ def get_s3_resource():
     return s3_resource
 
 
+def _get_bucket_name():
+    return current_app.config["CSV_UPLOAD_BUCKET"]["bucket"]
+
+
 def list_s3_objects():
 
-    bucket_name = current_app.config["CSV_UPLOAD_BUCKET"]["bucket"]
+    bucket_name = _get_bucket_name()
     s3_client = get_s3_client()
     # Our reports only support 7 days, but pull 8 days to avoid
     # any edge cases

tests/app/aws/test_s3.py

@@ -1,21 +1,32 @@
 import os
 from datetime import timedelta
 from os import getenv
+from unittest.mock import ANY, MagicMock, Mock, call, patch
 
+import botocore
 import pytest
 from botocore.exceptions import ClientError
 
 from app.aws.s3 import (
     cleanup_old_s3_objects,
+    download_from_s3,
     file_exists,
+    get_job_and_metadata_from_s3,
     get_job_from_s3,
     get_job_id_from_s3_object_key,
     get_personalisation_from_s3,
     get_phone_number_from_s3,
+    get_s3_client,
     get_s3_file,
+    get_s3_files,
+    get_s3_object,
+    get_s3_resource,
+    list_s3_objects,
+    read_s3_file,
     remove_csv_object,
     remove_s3_object,
 )
+from app.clients import AWS_CLIENT_CONFIG
 from app.utils import utc_now
 from notifications_utils import aware_utcnow
 
@@ -59,6 +70,110 @@ def test_cleanup_old_s3_objects(mocker):
     mock_remove_csv_object.assert_called_once_with("A")
 
 
+def test_read_s3_file_success(mocker):
+    mock_s3res = MagicMock()
+    mock_extract_personalisation = mocker.patch("app.aws.s3.extract_personalisation")
+    mock_extract_phones = mocker.patch("app.aws.s3.extract_phones")
+    mock_set_job_cache = mocker.patch("app.aws.s3.set_job_cache")
+    mock_get_job_id = mocker.patch("app.aws.s3.get_job_id_from_s3_object_key")
+    bucket_name = "test_bucket"
+    object_key = "test_object_key"
+    job_id = "12345"
+    file_content = "some file content"
+    mock_get_job_id.return_value = job_id
+    mock_s3_object = MagicMock()
+    mock_s3_object.get.return_value = {
+        "Body": MagicMock(read=MagicMock(return_value=file_content.encode("utf-8")))
+    }
+    mock_s3res.Object.return_value = mock_s3_object
+    mock_extract_phones.return_value = ["1234567890"]
+    mock_extract_personalisation.return_value = {"name": "John Doe"}
+
+    global job_cache
+    job_cache = {}
+
+    read_s3_file(bucket_name, object_key, mock_s3res)
+    mock_get_job_id.assert_called_once_with(object_key)
+    mock_s3res.Object.assert_called_once_with(bucket_name, object_key)
+    expected_calls = [
+        call(ANY, job_id, file_content),
+        call(ANY, f"{job_id}_phones", ["1234567890"]),
+        call(ANY, f"{job_id}_personalisation", {"name": "John Doe"}),
+    ]
+    mock_set_job_cache.assert_has_calls(expected_calls, any_order=True)
+
+
+def test_download_from_s3_success(mocker):
+    mock_s3 = MagicMock()
+    mock_get_s3_client = mocker.patch("app.aws.s3.get_s3_client")
+    mock_current_app = mocker.patch("app.aws.s3.current_app")
+    mock_logger = mock_current_app.logger
+    mock_get_s3_client.return_value = mock_s3
+    bucket_name = "test_bucket"
+    s3_key = "test_key"
+    local_filename = "test_file"
+    access_key = "access_key"
+    region = "test_region"
+    download_from_s3(
+        bucket_name, s3_key, local_filename, access_key, "secret_key", region
+    )
+    mock_s3.download_file.assert_called_once_with(bucket_name, s3_key, local_filename)
+    mock_logger.info.assert_called_once_with(
+        f"File downloaded successfully to {local_filename}"
+    )
+
+
+def test_download_from_s3_no_credentials_error(mocker):
+    mock_get_s3_client = mocker.patch("app.aws.s3.get_s3_client")
+    mock_current_app = mocker.patch("app.aws.s3.current_app")
+    mock_logger = mock_current_app.logger
+    mock_s3 = MagicMock()
+    mock_s3.download_file.side_effect = botocore.exceptions.NoCredentialsError
+    mock_get_s3_client.return_value = mock_s3
+    try:
+        download_from_s3(
+            "test_bucket", "test_key", "test_file", "access_key", "secret_key", "region"
+        )
+    except Exception:
+        pass
+    mock_logger.exception.assert_called_once_with("Credentials not found")
+
+
+def test_download_from_s3_general_exception(mocker):
+    mock_get_s3_client = mocker.patch("app.aws.s3.get_s3_client")
+    mock_current_app = mocker.patch("app.aws.s3.current_app")
+    mock_logger = mock_current_app.logger
+    mock_s3 = MagicMock()
+    mock_s3.download_file.side_effect = Exception()
+    mock_get_s3_client.return_value = mock_s3
+    try:
+        download_from_s3(
+            "test_bucket", "test_key", "test_file", "access_key", "secret_key", "region"
+        )
+    except Exception:
+        pass
+    mock_logger.exception.assert_called_once()
+
+
+def test_list_s3_objects(mocker):
+    mocker.patch("app.aws.s3._get_bucket_name", return_value="Foo")
+    mock_s3_client = mocker.Mock()
+    mocker.patch("app.aws.s3.get_s3_client", return_value=mock_s3_client)
+    lastmod30 = aware_utcnow() - timedelta(days=30)
+    lastmod3 = aware_utcnow() - timedelta(days=3)
+
+    mock_s3_client.list_objects_v2.side_effect = [
+        {
+            "Contents": [
+                {"Key": "A", "LastModified": lastmod30},
+                {"Key": "B", "LastModified": lastmod3},
+            ]
+        }
+    ]
+    result = list_s3_objects()
+    assert list(result) == ["B"]
+
+
 def test_get_s3_file_makes_correct_call(notify_api, mocker):
     get_s3_mock = mocker.patch("app.aws.s3.get_s3_object")
     get_s3_file(
@@ -154,6 +269,15 @@ def test_get_job_from_s3_exponential_backoff_on_throttling(mocker):
     assert mock_get_object.call_count == 8
 
 
+def test_get_job_from_s3_exponential_backoff_on_random_exception(mocker):
+    # We try multiple times to retrieve the job, and if we can't we return None
+    mock_get_object = mocker.patch("app.aws.s3.get_s3_object", side_effect=Exception())
+    mocker.patch("app.aws.s3.file_exists", return_value=True)
+    job = get_job_from_s3("service_id", "job_id")
+    assert job is None
+    assert mock_get_object.call_count == 1
+
+
 def test_get_job_from_s3_exponential_backoff_file_not_found(mocker):
     mock_get_object = mocker.patch("app.aws.s3.get_s3_object", return_value=None)
     mocker.patch("app.aws.s3.file_exists", return_value=False)
@@ -254,3 +378,153 @@ def test_file_exists_false(notify_api, mocker):
         )
 
     get_s3_mock.assert_called_once()
+
+
+def test_get_s3_files_success(notify_api, mocker):
+    mock_current_app = mocker.patch("app.aws.s3.current_app")
+    mock_current_app.config = {"CSV_UPLOAD_BUCKET": {"bucket": "test-bucket"}}
+    mock_thread_pool_executor = mocker.patch("app.aws.s3.ThreadPoolExecutor")
+    mock_read_s3_file = mocker.patch("app.aws.s3.read_s3_file")
+    mock_list_s3_objects = mocker.patch("app.aws.s3.list_s3_objects")
+    mock_get_s3_resource = mocker.patch("app.aws.s3.get_s3_resource")
+    mock_list_s3_objects.return_value = ["file1.csv", "file2.csv"]
+    mock_s3_resource = MagicMock()
+    mock_get_s3_resource.return_value = mock_s3_resource
+    mock_executor = MagicMock()
+
+    def mock_map(func, iterable):
+        for item in iterable:
+            func(item)
+
+    mock_executor.map.side_effect = mock_map
+    mock_thread_pool_executor.return_value.__enter__.return_value = mock_executor
+
+    get_s3_files()
+
+    # mock_current_app.config.__getitem__.assert_called_once_with("CSV_UPLOAD_BUCKET")
+    mock_list_s3_objects.assert_called_once()
+    mock_thread_pool_executor.assert_called_once()
+
+    mock_executor.map.assert_called_once()
+
+    calls = [
+        (("test-bucket", "file1.csv", mock_s3_resource),),
+        (("test-bucket", "file2.csv", mock_s3_resource),),
+    ]
+
+    mock_read_s3_file.assert_has_calls(calls, any_order=True)
+
+    # mock_current_app.info.assert_any_call("job_cache length before regen: 0 #notify-admin-1200")
+
+    # mock_current_app.info.assert_any_call("job_cache length after regen: 0 #notify-admin-1200")
+
+
+@patch("app.aws.s3.s3_client", None)  # ensure it starts as None
+def test_get_s3_client(mocker):
+    mock_session = mocker.patch("app.aws.s3.Session")
+    mock_current_app = mocker.patch("app.aws.s3.current_app")
+    sa_key = "sec"
+    sa_key = f"{sa_key}ret_access_key"
+    mock_current_app.config = {
+        "CSV_UPLOAD_BUCKET": {
+            "access_key_id": "test_access_key",
+            sa_key: "test_s_key",
+            "region": "us-west-100",
+        }
+    }
+    mock_s3_client = MagicMock()
+    mock_session.return_value.client.return_value = mock_s3_client
+    result = get_s3_client()
+
+    mock_session.return_value.client.assert_called_once_with("s3")
+    assert result == mock_s3_client
+
+
+@patch("app.aws.s3.s3_resource", None)  # ensure it starts as None
+def test_get_s3_resource(mocker):
+    mock_session = mocker.patch("app.aws.s3.Session")
+    mock_current_app = mocker.patch("app.aws.s3.current_app")
+    sa_key = "sec"
+    sa_key = f"{sa_key}ret_access_key"
+
+    mock_current_app.config = {
+        "CSV_UPLOAD_BUCKET": {
+            "access_key_id": "test_access_key",
+            sa_key: "test_s_key",
+            "region": "us-west-100",
+        }
+    }
+    mock_s3_resource = MagicMock()
+    mock_session.return_value.resource.return_value = mock_s3_resource
+    result = get_s3_resource()
+
+    mock_session.return_value.resource.assert_called_once_with(
+        "s3", config=AWS_CLIENT_CONFIG
+    )
+    assert result == mock_s3_resource
+
+
+def test_get_job_and_metadata_from_s3(mocker):
+    mock_get_s3_object = mocker.patch("app.aws.s3.get_s3_object")
+    mock_get_job_location = mocker.patch("app.aws.s3.get_job_location")
+
+    mock_get_job_location.return_value = {"bucket_name", "new_key"}
+    mock_s3_object = MagicMock()
+    mock_s3_object.get.return_value = {
+        "Body": MagicMock(read=MagicMock(return_value=b"job data")),
+        "Metadata": {"key": "value"},
+    }
+    mock_get_s3_object.return_value = mock_s3_object
+    result = get_job_and_metadata_from_s3("service_id", "job_id")
+
+    mock_get_job_location.assert_called_once_with("service_id", "job_id")
+    # mock_get_s3_object.assert_called_once_with("bucket_name", "new_key")
+    assert result == ("job data", {"key": "value"})
+
+
+def test_get_job_and_metadata_from_s3_fallback_to_old_location(mocker):
+    mock_get_job_location = mocker.patch("app.aws.s3.get_job_location")
+    mock_get_old_job_location = mocker.patch("app.aws.s3.get_old_job_location")
+    mock_get_job_location.return_value = {"bucket_name", "new_key"}
+    mock_get_s3_object = mocker.patch("app.aws.s3.get_s3_object")
+    # mock_get_s3_object.side_effect = [ClientError({"Error": {}}, "GetObject"), mock_s3_object]
+    mock_get_old_job_location.return_value = {"bucket_name", "old_key"}
+    mock_s3_object = MagicMock()
+    mock_s3_object.get.return_value = {
+        "Body": MagicMock(read=MagicMock(return_value=b"old job data")),
+        "Metadata": {"old_key": "old_value"},
+    }
+    mock_get_s3_object.side_effect = [
+        ClientError({"Error": {}}, "GetObject"),
+        mock_s3_object,
+    ]
+    result = get_job_and_metadata_from_s3("service_id", "job_id")
+    mock_get_job_location.assert_called_once_with("service_id", "job_id")
+    mock_get_old_job_location.assert_called_once_with("service_id", "job_id")
+    # mock_get_s3_object.assert_any_call("bucket_name", "new_key")
+    # mock_get_s3_object.assert_any_call("bucket_name", "old_key")
+    assert result == ("old job data", {"old_key": "old_value"})
+
+
+def test_get_s3_object_client_error(mocker):
+    mock_get_s3_resource = mocker.patch("app.aws.s3.get_s3_resource")
+    mock_current_app = mocker.patch("app.aws.s3.current_app")
+    mock_logger = mock_current_app.logger
+    mock_s3 = Mock()
+    mock_s3.Object.side_effect = botocore.exceptions.ClientError(
+        error_response={"Error": {"Code": "404", "Message": "Not Found"}},
+        operation_name="GetObject",
+    )
+    mock_get_s3_resource.return_value = mock_s3
+
+    bucket_name = "test-bucket"
+    file_location = "nonexistent-file.txt"
+    access_key = "test-access-key"
+    skey = "skey"
+    region = "us-west-200"
+    result = get_s3_object(bucket_name, file_location, access_key, skey, region)
+    assert result is None
+    mock_s3.Object.assert_called_once_with(bucket_name, file_location)
+    mock_logger.exception.assert_called_once_with(
+        f"Can't retrieve S3 Object from {file_location}"
+    )