class AWSReportDownloaderTest(MasuTestCase):
"""Test Cases for the AWS S3 functions."""
fake = Faker()
@classmethod
def setUpClass(cls):
cls.fake_customer_name = CUSTOMER_NAME
cls.fake_report_name = REPORT
cls.fake_bucket_prefix = PREFIX
cls.fake_bucket_name = BUCKET
cls.selected_region = REGION
cls.auth_credential = fake_arn(service='iam', generate_account_id=True)
cls.manifest_accessor = ReportManifestDBAccessor()
@classmethod
def tearDownClass(cls):
cls.manifest_accessor.close_session()
@patch('masu.util.aws.common.get_assume_role_session',
return_value=FakeSession)
def setUp(self, fake_session):
os.makedirs(DATA_DIR, exist_ok=True)
self.report_downloader = ReportDownloader(self.fake_customer_name,
self.auth_credential,
self.fake_bucket_name, 'AWS',
1)
self.aws_report_downloader = AWSReportDownloader(
**{
'customer_name': self.fake_customer_name,
'auth_credential': self.auth_credential,
'bucket': self.fake_bucket_name,
'report_name': self.fake_report_name,
'provider_id': 1
})
def tearDown(self):
shutil.rmtree(DATA_DIR, ignore_errors=True)
manifests = self.manifest_accessor._get_db_obj_query().all()
for manifest in manifests:
self.manifest_accessor.delete(manifest)
self.manifest_accessor.commit()
@patch('masu.external.downloader.aws.aws_report_downloader.boto3.resource')
@patch(
'masu.external.downloader.aws.aws_report_downloader.AWSReportDownloader.download_file',
return_value=(
'mock_file_name',
None,
))
def test_download_bucket(self, mock_boto_resource, mock_download_file):
"""Test download bucket method."""
mock_resource = Mock()
mock_bucket = Mock()
mock_bucket.objects.all.return_value = []
mock_resource.Bucket.return_value = mock_bucket
mock_boto_resource = mock_resource
out = self.aws_report_downloader.download_bucket()
expected_files = []
self.assertEqual(out, expected_files)
@patch(
'masu.external.downloader.aws.aws_report_downloader.AWSReportDownloader.download_file',
side_effect=mock_download_file_error)
def test_download_report_missing_manifest(self, mock_download_file):
fake_report_date = self.fake.date_time().replace(day=1)
out = self.report_downloader.download_report(fake_report_date)
self.assertEqual(out, [])
@patch('masu.external.report_downloader.ReportStatsDBAccessor')
@patch('masu.util.aws.common.get_assume_role_session',
return_value=FakeSessionDownloadError)
def test_download_report_missing_bucket(self, mock_stats, fake_session):
mock_stats.return_value.__enter__ = Mock()
fake_report_date = self.fake.date_time().replace(day=1)
fake_report_date_str = fake_report_date.strftime('%Y%m%dT000000.000Z')
expected_assembly_id = '882083b7-ea62-4aab-aa6a-f0d08d65ee2b'
input_key = f'/koku/20180701-20180801/{expected_assembly_id}/koku-1.csv.gz'
mock_manifest = {
'assemblyId': expected_assembly_id,
'billingPeriod': {
'start': fake_report_date_str
},
'reportKeys': [input_key]
}
with patch.object(AWSReportDownloader,
'_get_manifest',
return_value=mock_manifest):
with self.assertRaises(AWSReportDownloaderError):
report_downloader = ReportDownloader(self.fake_customer_name,
self.auth_credential,
self.fake_bucket_name,
'AWS', 2)
AWSReportDownloader(
**{
'customer_name': self.fake_customer_name,
'auth_credential': self.auth_credential,
'bucket': self.fake_bucket_name,
'report_name': self.fake_report_name,
'provider_id': 2
})
report_downloader.download_report(fake_report_date)
@patch('masu.util.aws.common.get_assume_role_session',
return_value=FakeSession)
def test_missing_report_name(self, fake_session):
"""Test downloading a report with an invalid report name."""
auth_credential = fake_arn(service='iam', generate_account_id=True)
with self.assertRaises(MasuProviderError):
AWSReportDownloader(self.fake_customer_name, auth_credential,
's3_bucket', 'wrongreport')
@patch('masu.util.aws.common.get_assume_role_session',
return_value=FakeSession)
def test_download_default_report(self, fake_session):
# actual test
auth_credential = fake_arn(service='iam', generate_account_id=True)
downloader = AWSReportDownloader(self.fake_customer_name,
auth_credential,
self.fake_bucket_name)
self.assertEqual(downloader.report_name, self.fake_report_name)
@patch('masu.util.aws.common.get_assume_role_session',
return_value=FakeSessionNoReport)
@patch('masu.util.aws.common.get_cur_report_definitions', return_value=[])
def test_download_default_report_no_report_found(self, fake_session,
fake_report_list):
auth_credential = fake_arn(service='iam', generate_account_id=True)
with self.assertRaises(MasuProviderError):
AWSReportDownloader(self.fake_customer_name, auth_credential,
self.fake_bucket_name)
@patch('masu.external.downloader.aws.aws_report_downloader.shutil')
@patch('masu.util.aws.common.get_assume_role_session',
return_value=FakeSession)
def test_check_size_success(self, fake_session, fake_shutil):
fake_client = Mock()
fake_client.get_object.return_value = {
'ContentLength': 123456,
'Body': Mock()
}
fake_shutil.disk_usage.return_value = (10, 10, 4096 * 1024 * 1024)
auth_credential = fake_arn(service='iam', generate_account_id=True)
downloader = AWSReportDownloader(self.fake_customer_name,
auth_credential,
self.fake_bucket_name)
downloader.s3_client = fake_client
fakekey = self.fake.file_path(depth=random.randint(1, 5),
extension=random.choice(
['json', 'csv.gz']))
result = downloader._check_size(fakekey, check_inflate=False)
self.assertTrue(result)
@patch('masu.external.downloader.aws.aws_report_downloader.shutil')
@patch('masu.util.aws.common.get_assume_role_session',
return_value=FakeSession)
def test_check_size_fail_nospace(self, fake_session, fake_shutil):
fake_client = Mock()
fake_client.get_object.return_value = {
'ContentLength': 123456,
'Body': Mock()
}
fake_shutil.disk_usage.return_value = (10, 10, 10)
auth_credential = fake_arn(service='iam', generate_account_id=True)
downloader = AWSReportDownloader(self.fake_customer_name,
auth_credential,
self.fake_bucket_name)
downloader.s3_client = fake_client
fakekey = self.fake.file_path(depth=random.randint(1, 5),
extension=random.choice(
['json', 'csv.gz']))
result = downloader._check_size(fakekey, check_inflate=False)
self.assertFalse(result)
@patch('masu.util.aws.common.get_assume_role_session',
return_value=FakeSession)
def test_check_size_fail_nosize(self, fake_session):
fake_client = Mock()
fake_client.get_object.return_value = {}
auth_credential = fake_arn(service='iam', generate_account_id=True)
downloader = AWSReportDownloader(self.fake_customer_name,
auth_credential,
self.fake_bucket_name)
downloader.s3_client = fake_client
fakekey = self.fake.file_path(depth=random.randint(1, 5),
extension=random.choice(
['json', 'csv.gz']))
with self.assertRaises(AWSReportDownloaderError):
downloader._check_size(fakekey, check_inflate=False)
@patch('masu.external.downloader.aws.aws_report_downloader.shutil')
@patch('masu.util.aws.common.get_assume_role_session',
return_value=FakeSession)
def test_check_size_inflate_success(self, fake_session, fake_shutil):
fake_client = Mock()
fake_client.get_object.return_value = {
'ContentLength': 123456,
'Body': io.BytesIO(b'\xd2\x02\x96I')
}
fake_shutil.disk_usage.return_value = (10, 10, 4096 * 1024 * 1024)
auth_credential = fake_arn(service='iam', generate_account_id=True)
downloader = AWSReportDownloader(self.fake_customer_name,
auth_credential,
self.fake_bucket_name)
downloader.s3_client = fake_client
fakekey = self.fake.file_path(depth=random.randint(1, 5),
extension='csv.gz')
result = downloader._check_size(fakekey, check_inflate=True)
self.assertTrue(result)
@patch('masu.external.downloader.aws.aws_report_downloader.shutil')
@patch('masu.util.aws.common.get_assume_role_session',
return_value=FakeSession)
def test_check_size_inflate_fail(self, fake_session, fake_shutil):
fake_client = Mock()
fake_client.get_object.return_value = {
'ContentLength': 123456,
'Body': io.BytesIO(b'\xd2\x02\x96I')
}
fake_shutil.disk_usage.return_value = (10, 10, 1234567)
auth_credential = fake_arn(service='iam', generate_account_id=True)
downloader = AWSReportDownloader(self.fake_customer_name,
auth_credential,
self.fake_bucket_name)
downloader.s3_client = fake_client
fakekey = self.fake.file_path(depth=random.randint(1, 5),
extension='csv.gz')
result = downloader._check_size(fakekey, check_inflate=True)
self.assertFalse(result)
@patch('masu.external.downloader.aws.aws_report_downloader.shutil')
@patch('masu.util.aws.common.get_assume_role_session',
return_value=FakeSession)
def test_download_file_check_size_fail(self, fake_session, fake_shutil):
fake_client = Mock()
fake_client.get_object.return_value = {
'ContentLength': 123456,
'Body': io.BytesIO(b'\xd2\x02\x96I')
}
fake_shutil.disk_usage.return_value = (10, 10, 1234567)
auth_credential = fake_arn(service='iam', generate_account_id=True)
downloader = AWSReportDownloader(self.fake_customer_name,
auth_credential,
self.fake_bucket_name)
downloader.s3_client = fake_client
fakekey = self.fake.file_path(depth=random.randint(1, 5),
extension='csv.gz')
with self.assertRaises(AWSReportDownloaderError):
downloader.download_file(fakekey)
@patch('masu.util.aws.common.get_assume_role_session',
return_value=FakeSession)
def test_download_file_raise_downloader_err(self, fake_session):
fake_response = {'Error': {'Code': self.fake.word()}}
fake_client = Mock()
fake_client.get_object.side_effect = ClientError(
fake_response, 'masu-test')
auth_credential = fake_arn(service='iam', generate_account_id=True)
downloader = AWSReportDownloader(self.fake_customer_name,
auth_credential,
self.fake_bucket_name)
downloader.s3_client = fake_client
with self.assertRaises(AWSReportDownloaderError):
downloader.download_file(self.fake.file_path())
@patch('masu.util.aws.common.get_assume_role_session',
return_value=FakeSession)
def test_download_file_raise_nofile_err(self, fake_session):
fake_response = {'Error': {'Code': 'NoSuchKey'}}
fake_client = Mock()
fake_client.get_object.side_effect = ClientError(
fake_response, 'masu-test')
auth_credential = fake_arn(service='iam', generate_account_id=True)
downloader = AWSReportDownloader(self.fake_customer_name,
auth_credential,
self.fake_bucket_name)
downloader.s3_client = fake_client
with self.assertRaises(AWSReportDownloaderNoFileError):
downloader.download_file(self.fake.file_path())
class AWSReportDownloaderTest(MasuTestCase):
"""Test Cases for the AWS S3 functions."""
fake = Faker()
@classmethod
def setUpClass(cls):
super().setUpClass()
cls.fake_customer_name = CUSTOMER_NAME
cls.fake_report_name = REPORT
cls.fake_bucket_prefix = PREFIX
cls.fake_bucket_name = BUCKET
cls.selected_region = REGION
cls.auth_credential = fake_arn(service='iam', generate_account_id=True)
cls.manifest_accessor = ReportManifestDBAccessor()
@patch('masu.util.aws.common.get_assume_role_session',
return_value=FakeSession)
def setUp(self, fake_session):
super().setUp()
os.makedirs(DATA_DIR, exist_ok=True)
self.mock_task = Mock(
request=Mock(id=str(self.fake.uuid4()), return_value={}))
self.report_downloader = ReportDownloader(
task=self.mock_task,
customer_name=self.fake_customer_name,
access_credential=self.auth_credential,
report_source=self.fake_bucket_name,
provider_type='AWS',
provider_uuid=self.aws_provider_uuid,
)
self.aws_report_downloader = AWSReportDownloader(
**{
'task': self.mock_task,
'customer_name': self.fake_customer_name,
'auth_credential': self.auth_credential,
'bucket': self.fake_bucket_name,
'report_name': self.fake_report_name,
'provider_uuid': self.aws_provider_uuid,
})
def tearDown(self):
shutil.rmtree(DATA_DIR, ignore_errors=True)
@patch('masu.external.downloader.aws.aws_report_downloader.boto3.resource')
@patch(
'masu.external.downloader.aws.aws_report_downloader.AWSReportDownloader.download_file',
return_value=('mock_file_name', None),
)
def test_download_bucket(self, mock_boto_resource, mock_download_file):
"""Test download bucket method."""
mock_resource = Mock()
mock_bucket = Mock()
mock_bucket.objects.all.return_value = []
mock_resource.Bucket.return_value = mock_bucket
mock_boto_resource = mock_resource
out = self.aws_report_downloader.download_bucket()
expected_files = []
self.assertEqual(out, expected_files)
@patch(
'masu.external.downloader.aws.aws_report_downloader.AWSReportDownloader.download_file',
side_effect=mock_download_file_error,
)
def test_download_report_missing_manifest(self, mock_download_file):
fake_report_date = self.fake.date_time().replace(day=1)
out = self.report_downloader.download_report(fake_report_date)
self.assertEqual(out, [])
@patch('masu.external.report_downloader.ReportStatsDBAccessor')
@patch(
'masu.util.aws.common.get_assume_role_session',
return_value=FakeSessionDownloadError,
)
def test_download_report_missing_bucket(self, mock_stats, fake_session):
mock_stats.return_value.__enter__ = Mock()
fake_report_date = self.fake.date_time().replace(day=1)
fake_report_date_str = fake_report_date.strftime('%Y%m%dT000000.000Z')
expected_assembly_id = '882083b7-ea62-4aab-aa6a-f0d08d65ee2b'
input_key = f'/koku/20180701-20180801/{expected_assembly_id}/koku-1.csv.gz'
mock_manifest = {
'assemblyId': expected_assembly_id,
'billingPeriod': {
'start': fake_report_date_str
},
'reportKeys': [input_key],
}
with patch.object(AWSReportDownloader,
'_get_manifest',
return_value=('', mock_manifest)):
with self.assertRaises(AWSReportDownloaderError):
report_downloader = ReportDownloader(
task=self.mock_task,
customer_name=self.fake_customer_name,
access_credential=self.auth_credential,
report_source=self.fake_bucket_name,
provider_type='AWS',
provider_uuid=self.aws_provider_uuid,
)
AWSReportDownloader(
**{
'task': self.mock_task,
'customer_name': self.fake_customer_name,
'auth_credential': self.auth_credential,
'bucket': self.fake_bucket_name,
'report_name': self.fake_report_name,
'provider_uuid': self.aws_provider_uuid,
})
report_downloader.download_report(fake_report_date)
@patch('masu.util.aws.common.get_assume_role_session',
return_value=FakeSession)
def test_missing_report_name(self, fake_session):
"""Test downloading a report with an invalid report name."""
auth_credential = fake_arn(service='iam', generate_account_id=True)
with self.assertRaises(MasuProviderError):
AWSReportDownloader(self.mock_task, self.fake_customer_name,
auth_credential, 's3_bucket', 'wrongreport')
@patch('masu.util.aws.common.get_assume_role_session',
return_value=FakeSession)
def test_download_default_report(self, fake_session):
# actual test
auth_credential = fake_arn(service='iam', generate_account_id=True)
downloader = AWSReportDownloader(self.mock_task,
self.fake_customer_name,
auth_credential,
self.fake_bucket_name)
self.assertEqual(downloader.report_name, self.fake_report_name)
@patch('masu.util.aws.common.get_assume_role_session',
return_value=FakeSessionNoReport)
@patch('masu.util.aws.common.get_cur_report_definitions', return_value=[])
def test_download_default_report_no_report_found(self, fake_session,
fake_report_list):
auth_credential = fake_arn(service='iam', generate_account_id=True)
with self.assertRaises(MasuProviderError):
AWSReportDownloader(self.mock_task, self.fake_customer_name,
auth_credential, self.fake_bucket_name)
@patch('masu.external.downloader.aws.aws_report_downloader.shutil')
@patch('masu.util.aws.common.get_assume_role_session',
return_value=FakeSession)
def test_check_size_success(self, fake_session, fake_shutil):
fake_client = Mock()
fake_client.get_object.return_value = {
'ContentLength': 123456,
'Body': Mock()
}
fake_shutil.disk_usage.return_value = (10, 10, 4096 * 1024 * 1024)
auth_credential = fake_arn(service='iam', generate_account_id=True)
downloader = AWSReportDownloader(self.mock_task,
self.fake_customer_name,
auth_credential,
self.fake_bucket_name)
downloader.s3_client = fake_client
fakekey = self.fake.file_path(depth=random.randint(1, 5),
extension=random.choice(
['json', 'csv.gz']))
result = downloader._check_size(fakekey, check_inflate=False)
self.assertTrue(result)
@patch('masu.external.downloader.aws.aws_report_downloader.shutil')
@patch('masu.util.aws.common.get_assume_role_session',
return_value=FakeSession)
def test_check_size_fail_nospace(self, fake_session, fake_shutil):
fake_client = Mock()
fake_client.get_object.return_value = {
'ContentLength': 123456,
'Body': Mock()
}
fake_shutil.disk_usage.return_value = (10, 10, 10)
auth_credential = fake_arn(service='iam', generate_account_id=True)
downloader = AWSReportDownloader(self.mock_task,
self.fake_customer_name,
auth_credential,
self.fake_bucket_name)
downloader.s3_client = fake_client
fakekey = self.fake.file_path(depth=random.randint(1, 5),
extension=random.choice(
['json', 'csv.gz']))
result = downloader._check_size(fakekey, check_inflate=False)
self.assertFalse(result)
@patch('masu.util.aws.common.get_assume_role_session',
return_value=FakeSession)
def test_check_size_fail_nosize(self, fake_session):
fake_client = Mock()
fake_client.get_object.return_value = {}
auth_credential = fake_arn(service='iam', generate_account_id=True)
downloader = AWSReportDownloader(self.mock_task,
self.fake_customer_name,
auth_credential,
self.fake_bucket_name)
downloader.s3_client = fake_client
fakekey = self.fake.file_path(depth=random.randint(1, 5),
extension=random.choice(
['json', 'csv.gz']))
with self.assertRaises(AWSReportDownloaderError):
downloader._check_size(fakekey, check_inflate=False)
@patch('masu.external.downloader.aws.aws_report_downloader.shutil')
@patch('masu.util.aws.common.get_assume_role_session',
return_value=FakeSession)
def test_check_size_inflate_success(self, fake_session, fake_shutil):
fake_client = Mock()
fake_client.get_object.return_value = {
'ContentLength': 123456,
'Body': io.BytesIO(b'\xd2\x02\x96I'),
}
fake_shutil.disk_usage.return_value = (10, 10, 4096 * 1024 * 1024)
auth_credential = fake_arn(service='iam', generate_account_id=True)
downloader = AWSReportDownloader(self.mock_task,
self.fake_customer_name,
auth_credential,
self.fake_bucket_name)
downloader.s3_client = fake_client
fakekey = self.fake.file_path(depth=random.randint(1, 5),
extension='csv.gz')
result = downloader._check_size(fakekey, check_inflate=True)
self.assertTrue(result)
@patch('masu.external.downloader.aws.aws_report_downloader.shutil')
@patch('masu.util.aws.common.get_assume_role_session',
return_value=FakeSession)
def test_check_size_inflate_fail(self, fake_session, fake_shutil):
fake_client = Mock()
fake_client.get_object.return_value = {
'ContentLength': 123456,
'Body': io.BytesIO(b'\xd2\x02\x96I'),
}
fake_shutil.disk_usage.return_value = (10, 10, 1234567)
auth_credential = fake_arn(service='iam', generate_account_id=True)
downloader = AWSReportDownloader(self.mock_task,
self.fake_customer_name,
auth_credential,
self.fake_bucket_name)
downloader.s3_client = fake_client
fakekey = self.fake.file_path(depth=random.randint(1, 5),
extension='csv.gz')
result = downloader._check_size(fakekey, check_inflate=True)
self.assertFalse(result)
@patch('masu.external.downloader.aws.aws_report_downloader.shutil')
@patch('masu.util.aws.common.get_assume_role_session',
return_value=FakeSession)
def test_download_file_check_size_fail(self, fake_session, fake_shutil):
fake_client = Mock()
fake_client.get_object.return_value = {
'ContentLength': 123456,
'Body': io.BytesIO(b'\xd2\x02\x96I'),
}
fake_shutil.disk_usage.return_value = (10, 10, 1234567)
auth_credential = fake_arn(service='iam', generate_account_id=True)
downloader = AWSReportDownloader(self.mock_task,
self.fake_customer_name,
auth_credential,
self.fake_bucket_name)
downloader.s3_client = fake_client
fakekey = self.fake.file_path(depth=random.randint(1, 5),
extension='csv.gz')
with self.assertRaises(AWSReportDownloaderError):
downloader.download_file(fakekey)
@patch('masu.util.aws.common.get_assume_role_session',
return_value=FakeSession)
def test_download_file_raise_downloader_err(self, fake_session):
fake_response = {'Error': {'Code': self.fake.word()}}
fake_client = Mock()
fake_client.get_object.side_effect = ClientError(
fake_response, 'masu-test')
auth_credential = fake_arn(service='iam', generate_account_id=True)
downloader = AWSReportDownloader(self.mock_task,
self.fake_customer_name,
auth_credential,
self.fake_bucket_name)
downloader.s3_client = fake_client
with self.assertRaises(AWSReportDownloaderError):
downloader.download_file(self.fake.file_path())
@patch('masu.util.aws.common.get_assume_role_session',
return_value=FakeSession)
def test_download_file_raise_nofile_err(self, fake_session):
fake_response = {'Error': {'Code': 'NoSuchKey'}}
fake_client = Mock()
fake_client.get_object.side_effect = ClientError(
fake_response, 'masu-test')
auth_credential = fake_arn(service='iam', generate_account_id=True)
downloader = AWSReportDownloader(self.mock_task,
self.fake_customer_name,
auth_credential,
self.fake_bucket_name)
downloader.s3_client = fake_client
with self.assertRaises(AWSReportDownloaderNoFileError):
downloader.download_file(self.fake.file_path())
@patch(
'masu.external.downloader.aws.aws_report_downloader.AWSReportDownloader.check_if_manifest_should_be_downloaded'
)
@patch(
'masu.external.downloader.aws.aws_report_downloader.AWSReportDownloader._remove_manifest_file'
)
@patch(
'masu.external.downloader.aws.aws_report_downloader.AWSReportDownloader._get_manifest'
)
@patch('masu.util.aws.common.get_assume_role_session',
return_value=FakeSession)
def test_get_report_context_for_date_should_download(
self, mock_session, mock_manifest, mock_delete, mock_check):
"""Test that data is returned on the reports to process."""
current_month = DateAccessor().today().replace(day=1,
second=1,
microsecond=1)
auth_credential = fake_arn(service='iam', generate_account_id=True)
downloader = AWSReportDownloader(self.mock_task,
self.fake_customer_name,
auth_credential,
self.fake_bucket_name,
provider_uuid=self.aws_provider_uuid)
start_str = current_month.strftime(downloader.manifest_date_format)
assembly_id = '1234'
compression = downloader.report.get('Compression')
report_keys = ['file1', 'file2']
mock_manifest.return_value = ('', {
'assemblyId': assembly_id,
'Compression': compression,
'reportKeys': report_keys,
'billingPeriod': {
'start': start_str
}
})
mock_check.return_value = True
expected = {
'manifest_id': None,
'assembly_id': assembly_id,
'compression': compression,
'files': report_keys
}
result = downloader.get_report_context_for_date(current_month)
with ReportManifestDBAccessor() as manifest_accessor:
manifest_entry = manifest_accessor.get_manifest(
assembly_id, self.aws_provider_uuid)
expected['manifest_id'] = manifest_entry.id
self.assertIsInstance(result, dict)
for key, value in result.items():
self.assertIn(key, expected)
self.assertEqual(value, expected.get(key))
@patch(
'masu.external.downloader.aws.aws_report_downloader.AWSReportDownloader.check_if_manifest_should_be_downloaded'
)
@patch(
'masu.external.downloader.aws.aws_report_downloader.AWSReportDownloader._remove_manifest_file'
)
@patch(
'masu.external.downloader.aws.aws_report_downloader.AWSReportDownloader._get_manifest'
)
@patch('masu.util.aws.common.get_assume_role_session',
return_value=FakeSession)
def test_get_report_context_for_date_should_not_download(
self, mock_session, mock_manifest, mock_delete, mock_check):
"""Test that no data is returned when we don't want to process."""
current_month = DateAccessor().today().replace(day=1,
second=1,
microsecond=1)
auth_credential = fake_arn(service='iam', generate_account_id=True)
downloader = AWSReportDownloader(self.mock_task,
self.fake_customer_name,
auth_credential,
self.fake_bucket_name)
start_str = current_month.strftime(downloader.manifest_date_format)
assembly_id = '1234'
compression = downloader.report.get('Compression')
report_keys = ['file1', 'file2']
mock_manifest.return_value = ('', {
'assemblyId': assembly_id,
'Compression': compression,
'reportKeys': report_keys,
'billingPeriod': {
'start': start_str
}
})
mock_check.return_value = False
expected = {}
result = downloader.get_report_context_for_date(current_month)
self.assertEqual(result, expected)
def test_remove_manifest_file(self):
manifest_file = f'{DATA_DIR}/test_manifest.json'
with open(manifest_file, 'w') as f:
f.write('Test')
self.assertTrue(os.path.isfile(manifest_file))
self.aws_report_downloader._remove_manifest_file(manifest_file)
self.assertFalse(os.path.isfile(manifest_file))
class AWSReportDownloaderTest(MasuTestCase):
"""Test Cases for the AWS S3 functions."""
fake = Faker()
@classmethod
def setUpClass(cls):
"""Set up shared class variables."""
super().setUpClass()
cls.fake_customer_name = CUSTOMER_NAME
cls.fake_report_name = REPORT
cls.fake_bucket_prefix = PREFIX
cls.fake_bucket_name = BUCKET
cls.selected_region = REGION
cls.auth_credential = fake_arn(service="iam", generate_account_id=True)
cls.manifest_accessor = ReportManifestDBAccessor()
@patch("masu.util.aws.common.get_assume_role_session", return_value=FakeSession)
def setUp(self, fake_session):
"""Set up shared variables."""
super().setUp()
os.makedirs(DATA_DIR, exist_ok=True)
self.mock_task = Mock(request=Mock(id=str(self.fake.uuid4()), return_value={}))
self.report_downloader = ReportDownloader(
task=self.mock_task,
customer_name=self.fake_customer_name,
access_credential=self.auth_credential,
report_source=self.fake_bucket_name,
provider_type=Provider.PROVIDER_AWS,
provider_uuid=self.aws_provider_uuid,
)
self.aws_report_downloader = AWSReportDownloader(
**{
"task": self.mock_task,
"customer_name": self.fake_customer_name,
"auth_credential": self.auth_credential,
"bucket": self.fake_bucket_name,
"report_name": self.fake_report_name,
"provider_uuid": self.aws_provider_uuid,
}
)
def tearDown(self):
"""Remove test generated data."""
shutil.rmtree(DATA_DIR, ignore_errors=True)
@patch("masu.external.downloader.aws.aws_report_downloader.boto3.resource")
@patch(
"masu.external.downloader.aws.aws_report_downloader.AWSReportDownloader.download_file",
return_value=("mock_file_name", None),
)
def test_download_bucket(self, mock_boto_resource, mock_download_file):
"""Test download bucket method."""
mock_resource = Mock()
mock_bucket = Mock()
mock_bucket.objects.all.return_value = []
mock_resource.Bucket.return_value = mock_bucket
out = self.aws_report_downloader.download_bucket()
expected_files = []
self.assertEqual(out, expected_files)
@patch(
"masu.external.downloader.aws.aws_report_downloader.AWSReportDownloader.download_file",
side_effect=mock_download_file_error,
)
def test_download_report_missing_manifest(self, mock_download_file):
"""Test download fails when manifest is missing."""
fake_report_date = self.fake.date_time().replace(day=1)
out = self.report_downloader.download_report(fake_report_date)
self.assertEqual(out, [])
@patch("masu.external.report_downloader.ReportStatsDBAccessor")
@patch("masu.util.aws.common.get_assume_role_session", return_value=FakeSessionDownloadError)
def test_download_report_missing_bucket(self, mock_stats, fake_session):
"""Test download fails when bucket is missing."""
mock_stats.return_value.__enter__ = Mock()
fake_report_date = self.fake.date_time().replace(day=1)
fake_report_date_str = fake_report_date.strftime("%Y%m%dT000000.000Z")
expected_assembly_id = "882083b7-ea62-4aab-aa6a-f0d08d65ee2b"
input_key = f"/koku/20180701-20180801/{expected_assembly_id}/koku-1.csv.gz"
mock_manifest = {
"assemblyId": expected_assembly_id,
"billingPeriod": {"start": fake_report_date_str},
"reportKeys": [input_key],
}
with patch.object(AWSReportDownloader, "_get_manifest", return_value=("", mock_manifest)):
with self.assertRaises(AWSReportDownloaderError):
report_downloader = ReportDownloader(
task=self.mock_task,
customer_name=self.fake_customer_name,
access_credential=self.auth_credential,
report_source=self.fake_bucket_name,
provider_type=Provider.PROVIDER_AWS,
provider_uuid=self.aws_provider_uuid,
)
AWSReportDownloader(
**{
"task": self.mock_task,
"customer_name": self.fake_customer_name,
"auth_credential": self.auth_credential,
"bucket": self.fake_bucket_name,
"report_name": self.fake_report_name,
"provider_uuid": self.aws_provider_uuid,
}
)
report_downloader.download_report(fake_report_date)
@patch("masu.util.aws.common.get_assume_role_session", return_value=FakeSession)
def test_missing_report_name(self, fake_session):
"""Test downloading a report with an invalid report name."""
auth_credential = fake_arn(service="iam", generate_account_id=True)
with self.assertRaises(MasuProviderError):
AWSReportDownloader(self.mock_task, self.fake_customer_name, auth_credential, "s3_bucket", "wrongreport")
@patch("masu.util.aws.common.get_assume_role_session", return_value=FakeSession)
def test_download_default_report(self, fake_session):
"""Test assume aws role works."""
# actual test
auth_credential = fake_arn(service="iam", generate_account_id=True)
downloader = AWSReportDownloader(
self.mock_task, self.fake_customer_name, auth_credential, self.fake_bucket_name
)
self.assertEqual(downloader.report_name, self.fake_report_name)
@patch("masu.util.aws.common.get_assume_role_session", return_value=FakeSessionNoReport)
@patch("masu.util.aws.common.get_cur_report_definitions", return_value=[])
def test_download_default_report_no_report_found(self, fake_session, fake_report_list):
"""Test download fails when no reports are found."""
auth_credential = fake_arn(service="iam", generate_account_id=True)
with self.assertRaises(MasuProviderError):
AWSReportDownloader(self.mock_task, self.fake_customer_name, auth_credential, self.fake_bucket_name)
@patch("masu.external.downloader.aws.aws_report_downloader.shutil")
@patch("masu.util.aws.common.get_assume_role_session", return_value=FakeSession)
def test_check_size_success(self, fake_session, fake_shutil):
"""Test _check_size is successful."""
fake_client = Mock()
fake_client.get_object.return_value = {"ContentLength": 123456, "Body": Mock()}
fake_shutil.disk_usage.return_value = (10, 10, 4096 * 1024 * 1024)
auth_credential = fake_arn(service="iam", generate_account_id=True)
downloader = AWSReportDownloader(
self.mock_task, self.fake_customer_name, auth_credential, self.fake_bucket_name
)
downloader.s3_client = fake_client
fakekey = self.fake.file_path(depth=random.randint(1, 5), extension=random.choice(["json", "csv.gz"]))
result = downloader._check_size(fakekey, check_inflate=False)
self.assertTrue(result)
@patch("masu.external.downloader.aws.aws_report_downloader.shutil")
@patch("masu.util.aws.common.get_assume_role_session", return_value=FakeSession)
def test_check_size_fail_nospace(self, fake_session, fake_shutil):
"""Test _check_size fails if there is no more space."""
fake_client = Mock()
fake_client.get_object.return_value = {"ContentLength": 123456, "Body": Mock()}
fake_shutil.disk_usage.return_value = (10, 10, 10)
auth_credential = fake_arn(service="iam", generate_account_id=True)
downloader = AWSReportDownloader(
self.mock_task, self.fake_customer_name, auth_credential, self.fake_bucket_name
)
downloader.s3_client = fake_client
fakekey = self.fake.file_path(depth=random.randint(1, 5), extension=random.choice(["json", "csv.gz"]))
result = downloader._check_size(fakekey, check_inflate=False)
self.assertFalse(result)
@patch("masu.util.aws.common.get_assume_role_session", return_value=FakeSession)
def test_check_size_fail_nosize(self, fake_session):
"""Test _check_size fails if there report has no size."""
fake_client = Mock()
fake_client.get_object.return_value = {}
auth_credential = fake_arn(service="iam", generate_account_id=True)
downloader = AWSReportDownloader(
self.mock_task, self.fake_customer_name, auth_credential, self.fake_bucket_name
)
downloader.s3_client = fake_client
fakekey = self.fake.file_path(depth=random.randint(1, 5), extension=random.choice(["json", "csv.gz"]))
with self.assertRaises(AWSReportDownloaderError):
downloader._check_size(fakekey, check_inflate=False)
@patch("masu.external.downloader.aws.aws_report_downloader.shutil")
@patch("masu.util.aws.common.get_assume_role_session", return_value=FakeSession)
def test_check_size_inflate_success(self, fake_session, fake_shutil):
"""Test _check_size inflation succeeds."""
fake_client = Mock()
fake_client.get_object.return_value = {"ContentLength": 123456, "Body": io.BytesIO(b"\xd2\x02\x96I")}
fake_shutil.disk_usage.return_value = (10, 10, 4096 * 1024 * 1024)
auth_credential = fake_arn(service="iam", generate_account_id=True)
downloader = AWSReportDownloader(
self.mock_task, self.fake_customer_name, auth_credential, self.fake_bucket_name
)
downloader.s3_client = fake_client
fakekey = self.fake.file_path(depth=random.randint(1, 5), extension="csv.gz")
result = downloader._check_size(fakekey, check_inflate=True)
self.assertTrue(result)
@patch("masu.external.downloader.aws.aws_report_downloader.shutil")
@patch("masu.util.aws.common.get_assume_role_session", return_value=FakeSession)
def test_check_size_inflate_fail(self, fake_session, fake_shutil):
"""Test _check_size fails when inflation fails."""
fake_client = Mock()
fake_client.get_object.return_value = {"ContentLength": 123456, "Body": io.BytesIO(b"\xd2\x02\x96I")}
fake_shutil.disk_usage.return_value = (10, 10, 1234567)
auth_credential = fake_arn(service="iam", generate_account_id=True)
downloader = AWSReportDownloader(
self.mock_task, self.fake_customer_name, auth_credential, self.fake_bucket_name
)
downloader.s3_client = fake_client
fakekey = self.fake.file_path(depth=random.randint(1, 5), extension="csv.gz")
result = downloader._check_size(fakekey, check_inflate=True)
self.assertFalse(result)
@patch("masu.external.downloader.aws.aws_report_downloader.shutil")
@patch("masu.util.aws.common.get_assume_role_session", return_value=FakeSession)
def test_download_file_check_size_fail(self, fake_session, fake_shutil):
"""Test _check_size fails when key is fake."""
fake_client = Mock()
fake_client.get_object.return_value = {"ContentLength": 123456, "Body": io.BytesIO(b"\xd2\x02\x96I")}
fake_shutil.disk_usage.return_value = (10, 10, 1234567)
auth_credential = fake_arn(service="iam", generate_account_id=True)
downloader = AWSReportDownloader(
self.mock_task, self.fake_customer_name, auth_credential, self.fake_bucket_name
)
downloader.s3_client = fake_client
fakekey = self.fake.file_path(depth=random.randint(1, 5), extension="csv.gz")
with self.assertRaises(AWSReportDownloaderError):
downloader.download_file(fakekey)
@patch("masu.util.aws.common.get_assume_role_session", return_value=FakeSession)
def test_download_file_raise_downloader_err(self, fake_session):
"""Test _check_size fails when there is a downloader error."""
fake_response = {"Error": {"Code": self.fake.word()}}
fake_client = Mock()
fake_client.get_object.side_effect = ClientError(fake_response, "masu-test")
auth_credential = fake_arn(service="iam", generate_account_id=True)
downloader = AWSReportDownloader(
self.mock_task, self.fake_customer_name, auth_credential, self.fake_bucket_name
)
downloader.s3_client = fake_client
with self.assertRaises(AWSReportDownloaderError):
downloader.download_file(self.fake.file_path())
@patch("masu.util.aws.common.get_assume_role_session", return_value=FakeSession)
def test_download_file_raise_nofile_err(self, fake_session):
"""Test that downloading a nonexistent file fails with AWSReportDownloaderNoFileError."""
fake_response = {"Error": {"Code": "NoSuchKey"}}
fake_client = Mock()
fake_client.get_object.side_effect = ClientError(fake_response, "masu-test")
auth_credential = fake_arn(service="iam", generate_account_id=True)
downloader = AWSReportDownloader(
self.mock_task, self.fake_customer_name, auth_credential, self.fake_bucket_name
)
downloader.s3_client = fake_client
with self.assertRaises(AWSReportDownloaderNoFileError):
downloader.download_file(self.fake.file_path())
@patch(
"masu.external.downloader.aws.aws_report_downloader.AWSReportDownloader.check_if_manifest_should_be_downloaded"
)
@patch("masu.external.downloader.aws.aws_report_downloader.AWSReportDownloader._remove_manifest_file")
@patch("masu.external.downloader.aws.aws_report_downloader.AWSReportDownloader._get_manifest")
@patch("masu.util.aws.common.get_assume_role_session", return_value=FakeSession)
def test_get_report_context_for_date_should_download(self, mock_session, mock_manifest, mock_delete, mock_check):
"""Test that data is returned on the reports to process."""
current_month = DateAccessor().today().replace(day=1, second=1, microsecond=1)
auth_credential = fake_arn(service="iam", generate_account_id=True)
downloader = AWSReportDownloader(
self.mock_task,
self.fake_customer_name,
auth_credential,
self.fake_bucket_name,
provider_uuid=self.aws_provider_uuid,
)
start_str = current_month.strftime(downloader.manifest_date_format)
assembly_id = "1234"
compression = downloader.report.get("Compression")
report_keys = ["file1", "file2"]
mock_manifest.return_value = (
"",
{
"assemblyId": assembly_id,
"Compression": compression,
"reportKeys": report_keys,
"billingPeriod": {"start": start_str},
},
)
mock_check.return_value = True
expected = {"manifest_id": None, "assembly_id": assembly_id, "compression": compression, "files": report_keys}
result = downloader.get_report_context_for_date(current_month)
with ReportManifestDBAccessor() as manifest_accessor:
manifest_entry = manifest_accessor.get_manifest(assembly_id, self.aws_provider_uuid)
expected["manifest_id"] = manifest_entry.id
self.assertIsInstance(result, dict)
for key, value in result.items():
self.assertIn(key, expected)
self.assertEqual(value, expected.get(key))
@patch(
"masu.external.downloader.aws.aws_report_downloader.AWSReportDownloader.check_if_manifest_should_be_downloaded"
)
@patch("masu.external.downloader.aws.aws_report_downloader.AWSReportDownloader._remove_manifest_file")
@patch("masu.external.downloader.aws.aws_report_downloader.AWSReportDownloader._get_manifest")
@patch("masu.util.aws.common.get_assume_role_session", return_value=FakeSession)
def test_get_report_context_for_date_should_not_download(
self, mock_session, mock_manifest, mock_delete, mock_check
):
"""Test that no data is returned when we don't want to process."""
current_month = DateAccessor().today().replace(day=1, second=1, microsecond=1)
auth_credential = fake_arn(service="iam", generate_account_id=True)
downloader = AWSReportDownloader(
self.mock_task, self.fake_customer_name, auth_credential, self.fake_bucket_name
)
start_str = current_month.strftime(downloader.manifest_date_format)
assembly_id = "1234"
compression = downloader.report.get("Compression")
report_keys = ["file1", "file2"]
mock_manifest.return_value = (
"",
{
"assemblyId": assembly_id,
"Compression": compression,
"reportKeys": report_keys,
"billingPeriod": {"start": start_str},
},
)
mock_check.return_value = False
expected = {}
result = downloader.get_report_context_for_date(current_month)
self.assertEqual(result, expected)
def test_remove_manifest_file(self):
"""Test that we remove the manifest file."""
manifest_file = f"{DATA_DIR}/test_manifest.json"
with open(manifest_file, "w") as f:
f.write("Test")
self.assertTrue(os.path.isfile(manifest_file))
self.aws_report_downloader._remove_manifest_file(manifest_file)
self.assertFalse(os.path.isfile(manifest_file))
def test_delete_manifest_file_warning(self):
"""Test that an INFO is logged when removing a manifest file that does not exist."""
with self.assertLogs(
logger="masu.external.downloader.aws.aws_report_downloader", level="INFO"
) as captured_logs:
# Disable log suppression
logging.disable(logging.NOTSET)
self.aws_report_downloader._remove_manifest_file("None")
self.assertTrue(
captured_logs.output[0].startswith("INFO:"),
msg="The log is expected to start with 'INFO:' but instead was: " + captured_logs.output[0],
)
self.assertTrue(
"Could not delete manifest file at" in captured_logs.output[0],
msg="""The log message is expected to contain
'Could not delete manifest file at' but instead was: """
+ captured_logs.output[0],
)
# Re-enable log suppression
logging.disable(logging.CRITICAL)
class AWSReportDownloaderTest(MasuTestCase):
"""Test Cases for the AWS S3 functions."""
fake = Faker()
@patch('masu.external.downloader.aws.utils.get_assume_role_session',
return_value=FakeSession)
def setUp(self, fake_session):
os.makedirs(DATA_DIR, exist_ok=True)
self.fake_customer_name = CUSTOMER_NAME
self.fake_report_name = REPORT
self.fake_bucket_name = BUCKET
self.fake_bucket_prefix = PREFIX
self.selected_region = REGION
auth_credential = fake_arn(service='iam', generate_account_id=True)
self.report_downloader = AWSReportDownloader(
**{
'customer_name': self.fake_customer_name,
'auth_credential': auth_credential,
'bucket': self.fake_bucket_name,
'report_name': self.fake_report_name
})
def tearDown(self):
shutil.rmtree(DATA_DIR, ignore_errors=True)
@mock_s3
def test_download_bucket(self):
fake_report_date = datetime.today().replace(day=1)
fake_report_end_date = fake_report_date + relativedelta(months=+1)
report_range = '{}-{}'.format(fake_report_date.strftime('%Y%m%d'),
fake_report_end_date.strftime('%Y%m%d'))
# Moto setup
conn = boto3.resource('s3', region_name=self.selected_region)
conn.create_bucket(Bucket=self.fake_bucket_name)
# push mocked csvs into Moto env
fake_csv_files = []
fake_csv_files_with_key = {}
for x in range(0, random.randint(2, 10)):
csv_filename = '{}.csv'.format('-'.join(
self.fake.words(random.randint(2, 5))))
fake_csv_files.append(csv_filename)
# mocked report file definition
fake_report_file = '{}/{}/{}/{}/{}'.format(self.fake_bucket_prefix,
self.fake_report_name,
report_range,
uuid.uuid4(),
csv_filename)
fake_csv_files_with_key[csv_filename] = fake_report_file
fake_csv_body = ','.join(self.fake.words(random.randint(5, 10)))
conn.Object(self.fake_bucket_name,
fake_report_file).put(Body=fake_csv_body)
key = conn.Object(self.fake_bucket_name, fake_report_file).get()
self.assertEqual(fake_csv_body, str(key['Body'].read(), 'utf-8'))
# mocked Manifest definition
fake_object = '{}/{}/{}/{}-Manifest.json'.format(
self.fake_bucket_prefix, self.fake_report_name, report_range,
self.fake_report_name)
fake_object_body = {'reportKeys': fake_csv_files}
# push mocked manifest into Moto env
conn.Object(self.fake_bucket_name,
fake_object).put(Body=json.dumps(fake_object_body))
key = conn.Object(self.fake_bucket_name, fake_object).get()
self.assertEqual(fake_object_body, json.load(key['Body']))
# actual test
out = self.report_downloader.download_bucket()
expected_files = []
for csv_filename in fake_csv_files:
report_key = fake_csv_files_with_key.get(csv_filename)
expected_assembly_id = utils.get_assembly_id_from_cur_key(
report_key)
expected_csv = '{}/{}/aws/{}/{}-{}'.format(DATA_DIR,
self.fake_customer_name,
self.fake_bucket_name,
expected_assembly_id,
csv_filename)
expected_files.append(expected_csv)
expected_manifest = '{}/{}/aws/{}/{}-Manifest.json'.format(
DATA_DIR, self.fake_customer_name, self.fake_bucket_name,
self.fake_report_name)
expected_files.append(expected_manifest)
self.assertEqual(sorted(out), sorted(expected_files))
@mock_s3
def test_download_current_report(self):
fake_report_date = datetime.today().replace(day=1)
fake_report_end_date = fake_report_date + relativedelta(months=+1)
report_range = '{}-{}'.format(fake_report_date.strftime('%Y%m%d'),
fake_report_end_date.strftime('%Y%m%d'))
# Moto setup
conn = boto3.resource('s3', region_name=self.selected_region)
conn.create_bucket(Bucket=self.fake_bucket_name)
# push mocked csvs into Moto env
fake_csv_files = {}
for x in range(0, random.randint(2, 10)):
csv_filename = '{}.csv'.format('-'.join(
self.fake.words(random.randint(2, 5))))
# mocked report file definition
fake_report_file = '{}/{}/{}/{}/{}'.format(self.fake_bucket_prefix,
self.fake_report_name,
report_range,
uuid.uuid4(),
csv_filename)
fake_csv_files[csv_filename] = fake_report_file
fake_csv_body = ','.join(self.fake.words(random.randint(5, 10)))
conn.Object(self.fake_bucket_name,
fake_report_file).put(Body=fake_csv_body)
key = conn.Object(self.fake_bucket_name, fake_report_file).get()
self.assertEqual(fake_csv_body, str(key['Body'].read(), 'utf-8'))
# mocked Manifest definition
selected_csv = random.choice(list(fake_csv_files.keys()))
fake_object = '{}/{}/{}/{}-Manifest.json'.format(
self.fake_bucket_prefix, self.fake_report_name, report_range,
self.fake_report_name)
fake_object_body = {'reportKeys': [fake_csv_files[selected_csv]]}
# push mocked manifest into Moto env
conn.Object(self.fake_bucket_name,
fake_object).put(Body=json.dumps(fake_object_body))
key = conn.Object(self.fake_bucket_name, fake_object).get()
self.assertEqual(fake_object_body, json.load(key['Body']))
# actual test
out = self.report_downloader.download_current_report()
files_list = []
for cur_dict in out:
files_list.append(cur_dict['file'])
self.assertIsNotNone(cur_dict['compression'])
report_key = fake_object_body.get('reportKeys').pop()
expected_assembly_id = utils.get_assembly_id_from_cur_key(report_key)
expected_csv = '{}/{}/aws/{}/{}-{}'.format(DATA_DIR,
self.fake_customer_name,
self.fake_bucket_name,
expected_assembly_id,
selected_csv)
self.assertEqual(files_list, [expected_csv])
# Verify etag is stored
for cur_dict in out:
cur_file = cur_dict['file']
file_name = cur_file.split('/')[-1]
stats_recorder = ReportStatsDBAccessor(file_name)
self.assertIsNotNone(stats_recorder.get_etag())
# Cleanup
stats_recorder.remove()
stats_recorder.commit()
stats_recorder2 = ReportStatsDBAccessor(file_name)
self.assertIsNone(stats_recorder2.get_etag())
stats_recorder.close_session()
stats_recorder2.close_session()
@mock_s3
def test_download_file(self):
fake_object = self.fake.word().lower()
conn = boto3.resource('s3', region_name=self.selected_region)
conn.create_bucket(Bucket=self.fake_bucket_name)
conn.Object(self.fake_bucket_name, fake_object).put(Body='test')
out, _ = self.report_downloader.download_file(fake_object)
self.assertEqual(
out, DATA_DIR + '/' + self.fake_customer_name + '/aws/' +
self.fake_bucket_name + '/' + fake_object)
@mock_s3
def test_download_file_missing_key(self):
fake_object = self.fake.word().lower()
conn = boto3.resource('s3', region_name=self.selected_region)
conn.create_bucket(Bucket=self.fake_bucket_name)
conn.Object(self.fake_bucket_name, fake_object).put(Body='test')
missing_key = 'missing' + fake_object
with self.assertRaises(AWSReportDownloaderError) as error:
self.report_downloader.download_file(missing_key)
expected_err = 'Unable to find {} in S3 Bucket: {}'.format(
missing_key, self.fake_bucket_name)
self.assertEqual(expected_err, str(error.exception))
@mock_s3
def test_download_file_other_error(self):
fake_object = self.fake.word().lower()
# No S3 bucket created
with self.assertRaises(AWSReportDownloaderError) as error:
self.report_downloader.download_file(fake_object)
self.assertTrue('NoSuchBucket' in str(error.exception))
@mock_s3
def test_download_report(self):
fake_report_date = self.fake.date_time().replace(day=1)
fake_report_end_date = fake_report_date + relativedelta(months=+1)
report_range = '{}-{}'.format(fake_report_date.strftime('%Y%m%d'),
fake_report_end_date.strftime('%Y%m%d'))
# mocked report file definition
fake_report_file = '{}/{}/{}/{}/{}.csv'.format(self.fake_bucket_prefix,
self.fake_report_name,
report_range,
uuid.uuid4(),
'mocked-report-file')
fake_report_file2 = '{}/{}/{}/{}/{}.csv'.format(
self.fake_bucket_prefix, self.fake_report_name, report_range,
uuid.uuid4(), 'mocked-report-file2')
# mocked Manifest definition
fake_object = '{}/{}/{}/{}-Manifest.json'.format(
self.fake_bucket_prefix, self.fake_report_name, report_range,
self.fake_report_name)
fake_object_body = {
'reportKeys': [fake_report_file, fake_report_file2]
}
# Moto setup
conn = boto3.resource('s3', region_name=self.selected_region)
conn.create_bucket(Bucket=self.fake_bucket_name)
# push mocked manifest into Moto env
conn.Object(self.fake_bucket_name,
fake_object).put(Body=json.dumps(fake_object_body))
key = conn.Object(self.fake_bucket_name, fake_object).get()
self.assertEqual(fake_object_body, json.load(key['Body']))
# push mocked csv into Moto env
fake_csv_body = ','.join(self.fake.words(random.randint(5, 10)))
conn.Object(self.fake_bucket_name,
fake_report_file).put(Body=fake_csv_body)
conn.Object(self.fake_bucket_name,
fake_report_file2).put(Body=fake_csv_body)
key = conn.Object(self.fake_bucket_name, fake_report_file).get()
self.assertEqual(fake_csv_body, str(key['Body'].read(), 'utf-8'))
# actual test. Run twice
for _ in range(2):
out = self.report_downloader.download_report(fake_report_date)
files_list = []
for cur_dict in out:
files_list.append(cur_dict['file'])
self.assertIsNotNone(cur_dict['compression'])
expected_paths = []
for report_key in fake_object_body.get('reportKeys'):
expected_assembly_id = utils.get_assembly_id_from_cur_key(
report_key)
expected_path_base = '{}/{}/{}/{}/{}-{}'
file_name = os.path.basename(report_key)
expected_path = expected_path_base.format(
DATA_DIR, self.fake_customer_name, 'aws',
self.fake_bucket_name, expected_assembly_id, file_name)
expected_paths.append(expected_path)
self.assertEqual(files_list, expected_paths)
@mock_s3
@patch('masu.external.downloader.aws.utils.get_assume_role_session',
return_value=FakeSession)
def test_missing_report_name(self, fake_session):
"""Test downloading a report with an invalid report name."""
auth_credential = fake_arn(service='iam', generate_account_id=True)
with self.assertRaises(MasuProviderError):
AWSReportDownloader(self.fake_customer_name, auth_credential,
's3_bucket', 'wrongreport')
@mock_s3
@patch('masu.external.downloader.aws.utils.get_assume_role_session',
return_value=FakeSession)
def test_download_default_report(self, fake_session):
fake_report_date = self.fake.date_time().replace(day=1)
fake_report_end_date = fake_report_date + relativedelta(months=+1)
report_range = '{}-{}'.format(fake_report_date.strftime('%Y%m%d'),
fake_report_end_date.strftime('%Y%m%d'))
# mocked report file definition
fake_report_file = '{}/{}/{}/{}/{}.csv'.format(self.fake_bucket_prefix,
self.fake_report_name,
report_range,
uuid.uuid4(),
'mocked-report-file')
# mocked Manifest definition
fake_object = '{}/{}/{}/{}-Manifest.json'.format(
self.fake_bucket_prefix, self.fake_report_name, report_range,
self.fake_report_name)
fake_object_body = {'reportKeys': [fake_report_file]}
# Moto setup
conn = boto3.resource('s3', region_name=self.selected_region)
conn.create_bucket(Bucket=self.fake_bucket_name)
# push mocked manifest into Moto env
conn.Object(self.fake_bucket_name,
fake_object).put(Body=json.dumps(fake_object_body))
key = conn.Object(self.fake_bucket_name, fake_object).get()
self.assertEqual(fake_object_body, json.load(key['Body']))
# push mocked csv into Moto env
fake_csv_body = ','.join(self.fake.words(random.randint(5, 10)))
conn.Object(self.fake_bucket_name,
fake_report_file).put(Body=fake_csv_body)
key = conn.Object(self.fake_bucket_name, fake_report_file).get()
self.assertEqual(fake_csv_body, str(key['Body'].read(), 'utf-8'))
# actual test
auth_credential = fake_arn(service='iam', generate_account_id=True)
downloader = AWSReportDownloader(self.fake_customer_name,
auth_credential,
self.fake_bucket_name)
self.assertEqual(downloader.report_name, self.fake_report_name)
@mock_s3
@patch('masu.external.downloader.aws.utils.get_assume_role_session',
return_value=FakeSessionNoReport)
@patch('masu.external.downloader.aws.utils.get_cur_report_definitions',
return_value=[])
def test_download_default_report_no_report_found(self, fake_session,
fake_report_list):
auth_credential = fake_arn(service='iam', generate_account_id=True)
with self.assertRaises(MasuProviderError):
AWSReportDownloader(self.fake_customer_name, auth_credential,
self.fake_bucket_name)