diff --git a/src/registrar/tests/test_reports.py b/src/registrar/tests/test_reports.py index e8796e1fa..23ac04139 100644 --- a/src/registrar/tests/test_reports.py +++ b/src/registrar/tests/test_reports.py @@ -1,4 +1,5 @@ import csv +import io from django.test import Client, RequestFactory, TestCase from io import StringIO from registrar.models.domain_information import DomainInformation @@ -11,7 +12,8 @@ from unittest.mock import MagicMock, call, mock_open, patch from api.views import get_current_federal, get_current_full from django.conf import settings from botocore.exceptions import ClientError -import boto3_mocking # type: ignore +import boto3_mocking +from registrar.utility.s3_bucket import S3ClientError, S3ClientErrorCodes # type: ignore class CsvReportsTest(TestCase): """Tests to determine if we are uploading our reports correctly""" @@ -124,75 +126,114 @@ class CsvReportsTest(TestCase): @boto3_mocking.patching def test_not_found_full_report(self): """Ensures that we get a not found when the report doesn't exist""" - def side_effect(fake): + def side_effect(Bucket, Key): raise ClientError({"Error": {"Code": "NoSuchKey", "Message": "No such key"}}, "get_object") mock_client = MagicMock() - mock_client_instance = mock_client.return_value mock_client.get_object.side_effect = side_effect - with patch('boto3.client', return_value=mock_client): - with self.assertRaises(ClientError) as context: - with boto3_mocking.clients.handler_for("s3", mock_client): - response = self.client.get("/api/v1/get-report/current-full") - expected_call = [ - call.get_object(Bucket=settings.AWS_S3_BUCKET_NAME, Key='current-full.csv') - ] - mock_client_instance.assert_has_calls(expected_call) - mock_client_instance.get_object.side_effect = Exception("An error occurred") - print("look") - print(response.content.decode()) - # Check that the response has status code 404 - self.assertEqual(response.status_code, 404) - # Check that the response body contains "File not found" - self.assertEqual(response.content.decode(), "File not found") + response = None + with boto3_mocking.clients.handler_for("s3", mock_client): + with patch('boto3.client', return_value=mock_client): + with self.assertRaises(S3ClientError) as context: + response = self.client.get("/api/v1/get-report/current-full") + # Check that the response has status code 500 + self.assertEqual(response.status_code, 500) + + # Check that we get the right error back from the page + self.assertEqual(context.exception.code, S3ClientErrorCodes.FILE_NOT_FOUND_ERROR) + @boto3_mocking.patching def test_not_found_federal_report(self): """Ensures that we get a not found when the report doesn't exist""" - response = self.client.get("/api/v1/get-report/current-federal") + def side_effect(Bucket, Key): + raise ClientError({"Error": {"Code": "NoSuchKey", "Message": "No such key"}}, "get_object") + mock_client = MagicMock() + mock_client.get_object.side_effect = side_effect + + with boto3_mocking.clients.handler_for("s3", mock_client): + with patch('boto3.client', return_value=mock_client): + with self.assertRaises(S3ClientError) as context: + response = self.client.get("/api/v1/get-report/current-federal") + # Check that the response has status code 500 + self.assertEqual(response.status_code, 500) - # Check that the response has status code 404 - self.assertEqual(response.status_code, 404) - # Check that the response body contains "File not found" - self.assertEqual(response.content.decode(), "File not found") + # Check that we get the right error back from the page + self.assertEqual(context.exception.code, S3ClientErrorCodes.FILE_NOT_FOUND_ERROR) @boto3_mocking.patching def test_load_federal_report(self): - """Tests the current-federal api link""" - if not boto3_mocking.patching_engaged(): - raise Exception("test123") - request = self.factory.get("/fake-path") - response = get_current_federal(request, file_path="registrar/tests/data/fake_current_federal.csv") + """Tests the get_current_federal api endpoint""" + self.maxDiff = None + mock_client = MagicMock() + mock_client_instance = mock_client.return_value + + with open("registrar/tests/data/fake_current_federal.csv", "r") as file: + file_content = file.read() + + # Mock a recieved file + mock_client_instance.get_object.return_value = { + 'Body': io.BytesIO(file_content.encode()) + } + with boto3_mocking.clients.handler_for("s3", mock_client): + request = self.factory.get("/fake-path") + response = get_current_federal(request) + + # Check that we are sending the correct calls. + # Ensures that we are decoding the file content recieved from AWS. + expected_call = [ + call.get_object(Bucket=settings.AWS_S3_BUCKET_NAME, Key='current-federal.csv') + ] + mock_client_instance.assert_has_calls(expected_call) + # Check that the response has status code 200 self.assertEqual(response.status_code, 200) - # Check that the response contains what we expect - file_content = b"".join(response.streaming_content).decode("utf-8") + # Check that the response contains what we expect expected_file_content = ( "Domain name,Domain type,Agency,Organization name,City,State,Security Contact Email\n" "cdomain1.gov,Federal - Executive,World War I Centennial Commission,,,,\n" "ddomain3.gov,Federal,Armed Forces Retirement Home,,,," - ) + ).encode() - self.assertEqual(file_content, expected_file_content) + self.assertEqual(expected_file_content, response.content) @boto3_mocking.patching def test_load_full_report(self): """Tests the current-federal api link""" - request = self.factory.get("/fake-path") - response = get_current_full(request, file_path="registrar/tests/data/fake_current_full.csv") + mock_client = MagicMock() + mock_client_instance = mock_client.return_value + + with open("registrar/tests/data/fake_current_full.csv", "r") as file: + file_content = file.read() + + # Mock a recieved file + mock_client_instance.get_object.return_value = { + 'Body': io.BytesIO(file_content.encode()) + } + with boto3_mocking.clients.handler_for("s3", mock_client): + request = self.factory.get("/fake-path") + response = get_current_federal(request) + + # Check that we are sending the correct calls. + # Ensures that we are decoding the file content recieved from AWS. + expected_call = [ + call.get_object(Bucket=settings.AWS_S3_BUCKET_NAME, Key='current-federal.csv') + ] + mock_client_instance.assert_has_calls(expected_call) + # Check that the response has status code 200 self.assertEqual(response.status_code, 200) + # Check that the response contains what we expect - file_content = b"".join(response.streaming_content).decode("utf-8") expected_file_content = ( "Domain name,Domain type,Agency,Organization name,City,State,Security Contact Email\n" "cdomain1.gov,Federal - Executive,World War I Centennial Commission,,,,\n" "ddomain3.gov,Federal,Armed Forces Retirement Home,,,,\n" "adomain2.gov,Interstate,,,,," - ) + ).encode() - self.assertEqual(file_content, expected_file_content) + self.assertEqual(expected_file_content, response.content) class ExportDataTest(TestCase): diff --git a/src/registrar/utility/s3_bucket.py b/src/registrar/utility/s3_bucket.py index 329a68c40..6370b0736 100644 --- a/src/registrar/utility/s3_bucket.py +++ b/src/registrar/utility/s3_bucket.py @@ -1,14 +1,45 @@ """Utilities for accessing an AWS S3 bucket""" +from enum import IntEnum import boto3 - +from botocore.exceptions import ClientError from django.conf import settings from django.template.loader import get_template +class S3ClientErrorCodes(IntEnum): + """Used for S3ClientError + Error code overview: + - 1 ACCESS_S3_CLIENT_ERROR + - 2 UPLOAD_FILE_ERROR + - 3 FILE_NOT_FOUND_ERROR + - 4 GET_FILE_ERROR + """ + + ACCESS_S3_CLIENT_ERROR = 1 + UPLOAD_FILE_ERROR = 2 + FILE_NOT_FOUND_ERROR = 3 + GET_FILE_ERROR = 4 class S3ClientError(RuntimeError): """Local error for handling all failures with boto3.client""" - pass + _error_mapping = { + S3ClientErrorCodes.ACCESS_S3_CLIENT_ERROR: "Failed to establish a connection with the storage service.", + S3ClientErrorCodes.UPLOAD_FILE_ERROR: "File upload to the storage service failed.", + S3ClientErrorCodes.FILE_NOT_FOUND_ERROR: "Requested file not found in the storage service.", + S3ClientErrorCodes.GET_FILE_ERROR: ( + "Retrieval of the requested file from " + "the storage service failed due to an unspecified error." + ), + } + + def __init__(self, *args, code=None, **kwargs): + super().__init__(*args, **kwargs) + self.code = code + if self.code in self._error_mapping: + self.message = self._error_mapping.get(self.code) + + def __str__(self): + return f"{self.message}" class S3ClientHelper: @@ -23,7 +54,7 @@ class S3ClientHelper: config=settings.BOTO_CONFIG, ) except Exception as exc: - raise S3ClientError("Could not access the S3 client") from exc + raise S3ClientError(code=S3ClientErrorCodes.ACCESS_S3_CLIENT_ERROR) from exc def get_bucket_name(self): """Gets the name of our S3 Bucket""" @@ -34,17 +65,20 @@ class S3ClientHelper: try: response = self.boto_client.upload_file(file_path, self.get_bucket_name(), file_name) except Exception as exc: - raise S3ClientError("Couldn't upload file") from exc + raise S3ClientError(code=S3ClientErrorCodes.UPLOAD_FILE_ERROR) from exc return response def get_file(self, file_name, decode_to_utf=False): """Gets a file to our S3 instance and returns the file content""" try: response = self.boto_client.get_object(Bucket=self.get_bucket_name(), Key=file_name) - except self.boto_client.exceptions.NoSuchKey as exc: - raise S3ClientError("File was not found") from exc + except ClientError as exc: + if exc.response['Error']['Code'] == 'NoSuchKey': + raise S3ClientError(code=S3ClientErrorCodes.FILE_NOT_FOUND_ERROR) from exc + else: + raise S3ClientError(code=S3ClientErrorCodes.GET_FILE_ERROR) from exc except Exception as exc: - raise S3ClientError("Couldn't get file, an unspecified error occured") from exc + raise S3ClientError(code=S3ClientErrorCodes.GET_FILE_ERROR) from exc file_content = response["Body"].read() if decode_to_utf: