def test_get_manifest_context_for_date(self, mock_session, mock_manifest,
mock_delete):
"""Test that the manifest is read."""
current_month = DateAccessor().today().replace(day=1,
second=1,
microsecond=1)
downloader = AWSReportDownloader(self.fake_customer_name,
self.credentials,
self.data_source,
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
},
},
DateAccessor().today(),
)
result = downloader.get_manifest_context_for_date(current_month)
self.assertEqual(result.get("assembly_id"), assembly_id)
self.assertEqual(result.get("compression"), compression)
self.assertIsNotNone(result.get("files"))
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 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,
cache_key=self.fake.word(),
)
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 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_get_manifest_context_for_date_no_manifest(self, mock_session, mock_manifest, mock_delete):
"""Test that the manifest is read."""
current_month = DateAccessor().today().replace(day=1, second=1, microsecond=1)
downloader = AWSReportDownloader(
self.fake_customer_name, self.credentials, self.data_source, provider_uuid=self.aws_provider_uuid
)
mock_manifest.return_value = ("", {"reportKeys": []}, DateAccessor().today())
result = downloader.get_manifest_context_for_date(current_month)
self.assertEqual(result, {})
def test_download_file_raise_access_denied_err(self, fake_session):
"""Test that downloading a fails when accessdenied occurs."""
fake_response = {"Error": {"Code": "AccessDenied"}}
fake_client = Mock()
fake_client.get_object.side_effect = ClientError(fake_response, "masu-test")
downloader = AWSReportDownloader(self.fake_customer_name, self.credentials, self.data_source)
downloader.s3_client = fake_client
with self.assertRaises(AWSReportDownloaderNoFileError):
downloader.download_file(self.fake.file_path())
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")
downloader = AWSReportDownloader(self.fake_customer_name, self.credentials, self.data_source)
downloader.s3_client = fake_client
with self.assertRaises(AWSReportDownloaderNoFileError):
downloader.download_file(self.fake.file_path())
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")
downloader = AWSReportDownloader(self.fake_customer_name, self.credentials, self.data_source)
downloader.s3_client = fake_client
with self.assertRaises(AWSReportDownloaderError):
downloader.download_file(self.fake.file_path())
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 = {}
downloader = AWSReportDownloader(self.fake_customer_name, self.credentials, self.data_source)
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)
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)
downloader = AWSReportDownloader(self.fake_customer_name, self.credentials, self.data_source)
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)
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)
downloader = AWSReportDownloader(self.fake_customer_name, self.credentials, self.data_source)
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)
def test_check_size_fail_unknown_error(self, fake_session):
"""Test _check_size fails if there report has no size."""
fake_response = {"Error": {"Code": "Unknown"}}
fake_client = Mock()
fake_client.get_object.side_effect = ClientError(fake_response, "masu-test")
downloader = AWSReportDownloader(self.fake_customer_name, self.credentials, self.data_source)
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)
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)
downloader = AWSReportDownloader(self.fake_customer_name, self.credentials, self.data_source)
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)
def test_download_file(self, fake_session, mock_copy, mock_remove,
mock_check_csvs, mock_mark_csvs, mock_check_size):
"""Test the download file method."""
mock_check_csvs.return_value = False
mock_check_size.return_value = True
downloader = AWSReportDownloader(self.fake_customer_name,
self.credentials, self.data_source)
downloader.download_file(self.fake.file_path(), manifest_id=1)
mock_check_csvs.assert_called()
mock_copy.assert_called()
mock_remove.assert_called()
mock_mark_csvs.assert_called()
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())
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())
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())
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)
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))
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))
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 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)
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)
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)
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)
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')
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)
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)
def _set_downloader(self):
"""
Create the report downloader object.
Downloader is specific to the provider's cloud service.
Args:
None
Returns:
(Object) : Some object that is a child of CURAccountsInterface
"""
if self.provider_type == AMAZON_WEB_SERVICES:
return AWSReportDownloader(customer_name=self.customer_name,
auth_credential=self.credential,
bucket=self.cur_source,
report_name=self.report_name)
elif self.provider_type == LOCAL_SERVICE_PROVIDER:
return LocalReportDownloader(customer_name=self.customer_name,
auth_credential=self.credential,
bucket=self.cur_source,
report_name=self.report_name)
return None
def test_init_with_demo_account(self):
"""Test init with the demo account."""
mock_task = Mock(
request=Mock(id=str(self.fake.uuid4()), return_value={}))
account_id = "123456"
arn = fake_arn(service="iam", generate_account_id=True)
credentials = {"role_arn": arn}
report_name = FAKE.word()
demo_accounts = {
account_id: {
arn: {
"report_name": report_name,
"report_prefix": FAKE.word()
}
}
}
with self.settings(DEMO_ACCOUNTS=demo_accounts):
with patch("masu.util.aws.common.get_assume_role_session"
) as mock_session:
AWSReportDownloader(
**{
"task": mock_task,
"customer_name": f"acct{account_id}",
"credentials": credentials,
"data_source": {
"bucket": FAKE.word()
},
"report_name": report_name,
"provider_uuid": self.aws_provider_uuid,
})
mock_session.assert_called_once()
