From ea289a6f592a8e6009b09fd8ca5083dc5442d3bb Mon Sep 17 00:00:00 2001 From: zandercymatics <141044360+zandercymatics@users.noreply.github.com> Date: Fri, 1 Dec 2023 09:06:27 -0700 Subject: [PATCH] Linting + test cases --- src/api/views.py | 3 +- .../generate_current_federal_report.py | 6 +-- .../commands/generate_current_full_report.py | 7 ++-- src/registrar/tests/test_reports.py | 42 +++++++++---------- src/registrar/utility/s3_bucket.py | 10 +++-- 5 files changed, 32 insertions(+), 36 deletions(-) diff --git a/src/api/views.py b/src/api/views.py index 60f11023e..f888ee6d4 100644 --- a/src/api/views.py +++ b/src/api/views.py @@ -1,8 +1,7 @@ """Internal API views""" -import os from django.apps import apps from django.views.decorators.http import require_http_methods -from django.http import FileResponse, HttpResponse, JsonResponse +from django.http import HttpResponse, JsonResponse from django.utils.safestring import mark_safe from registrar.templatetags.url_helpers import public_site_url diff --git a/src/registrar/management/commands/generate_current_federal_report.py b/src/registrar/management/commands/generate_current_federal_report.py index 6075d8ebb..478265ca0 100644 --- a/src/registrar/management/commands/generate_current_federal_report.py +++ b/src/registrar/management/commands/generate_current_federal_report.py @@ -39,11 +39,11 @@ class Command(BaseCommand): logger.info(f"Success! Created {file_name}") def generate_current_federal_report(self, directory, file_name, check_path): - """Creates a current-full.csv file under the specified directory, + """Creates a current-full.csv file under the specified directory, then uploads it to a AWS S3 bucket""" s3_client = S3ClientHelper() file_path = os.path.join(directory, file_name) - + # Generate a file locally for upload with open(file_path, "w") as file: csv_export.export_data_federal_to_csv(file) @@ -51,5 +51,5 @@ class Command(BaseCommand): if check_path and not os.path.exists(file_path): raise FileNotFoundError(f"Could not find newly created file at '{file_path}'") - # Upload this generated file for our S3 instance + # Upload this generated file for our S3 instance s3_client.upload_file(file_path, file_name) diff --git a/src/registrar/management/commands/generate_current_full_report.py b/src/registrar/management/commands/generate_current_full_report.py index 83be4f9b5..c41c697a1 100644 --- a/src/registrar/management/commands/generate_current_full_report.py +++ b/src/registrar/management/commands/generate_current_full_report.py @@ -39,11 +39,11 @@ class Command(BaseCommand): logger.info(f"Success! Created {file_name}") def generate_current_full_report(self, directory, file_name, check_path): - """Creates a current-full.csv file under the specified directory, + """Creates a current-full.csv file under the specified directory, then uploads it to a AWS S3 bucket""" s3_client = S3ClientHelper() file_path = os.path.join(directory, file_name) - + # Generate a file locally for upload with open(file_path, "w") as file: csv_export.export_data_full_to_csv(file) @@ -51,6 +51,5 @@ class Command(BaseCommand): if check_path and not os.path.exists(file_path): raise FileNotFoundError(f"Could not find newly created file at '{file_path}'") - # Upload this generated file for our S3 instance + # Upload this generated file for our S3 instance s3_client.upload_file(file_path, file_name) - diff --git a/src/registrar/tests/test_reports.py b/src/registrar/tests/test_reports.py index a06d6f1c2..266f7c799 100644 --- a/src/registrar/tests/test_reports.py +++ b/src/registrar/tests/test_reports.py @@ -15,6 +15,7 @@ from botocore.exceptions import ClientError 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""" @@ -106,18 +107,20 @@ class CsvReportsTest(TestCase): content = fake_open() # Now you can make assertions about how you expect 'file' to be used. content.write.assert_has_calls(expected_file_content) - + @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(Bucket, Key): raise ClientError({"Error": {"Code": "NoSuchKey", "Message": "No such key"}}, "get_object") + mock_client = MagicMock() mock_client.get_object.side_effect = side_effect response = None with boto3_mocking.clients.handler_for("s3", mock_client): - with patch('boto3.client', return_value=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 @@ -126,17 +129,18 @@ class CsvReportsTest(TestCase): # 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""" + 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 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 @@ -144,7 +148,7 @@ class CsvReportsTest(TestCase): # 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 get_current_federal api endpoint""" @@ -156,18 +160,14 @@ class CsvReportsTest(TestCase): file_content = file.read() # Mock a recieved file - mock_client_instance.get_object.return_value = { - 'Body': io.BytesIO(file_content.encode()) - } + 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. + + # 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') - ] + 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 @@ -192,18 +192,14 @@ class CsvReportsTest(TestCase): file_content = file.read() # Mock a recieved file - mock_client_instance.get_object.return_value = { - 'Body': io.BytesIO(file_content.encode()) - } + 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. + response = get_current_full(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') - ] + expected_call = [call.get_object(Bucket=settings.AWS_S3_BUCKET_NAME, Key="current-full.csv")] mock_client_instance.assert_has_calls(expected_call) # Check that the response has status code 200 diff --git a/src/registrar/utility/s3_bucket.py b/src/registrar/utility/s3_bucket.py index 6370b0736..162be37c9 100644 --- a/src/registrar/utility/s3_bucket.py +++ b/src/registrar/utility/s3_bucket.py @@ -4,7 +4,7 @@ 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 @@ -20,15 +20,16 @@ class S3ClientErrorCodes(IntEnum): FILE_NOT_FOUND_ERROR = 3 GET_FILE_ERROR = 4 + class S3ClientError(RuntimeError): """Local error for handling all failures with boto3.client""" + _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." + "Retrieval of the requested file from " "the storage service failed due to an unspecified error." ), } @@ -44,6 +45,7 @@ class S3ClientError(RuntimeError): class S3ClientHelper: """Helper class that simplifies S3 intialization""" + def __init__(self): try: self.boto_client = boto3.client( @@ -73,7 +75,7 @@ class S3ClientHelper: try: response = self.boto_client.get_object(Bucket=self.get_bucket_name(), Key=file_name) except ClientError as exc: - if exc.response['Error']['Code'] == 'NoSuchKey': + 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