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)