def test_upload_to_gcs_with_csv(self, mock_os, mock_tempfile,
mock_storage):
"""Test run()."""
fake_tmpname = 'tmp_name'
fake_output_name = 'abc'
mock_service_config = mock.MagicMock()
mock_service_config.get_notifier_config.return_value = {
'inventory': {
'gcs_summary': {
'enabled': True,
'data_format': 'csv',
'gcs_path': 'gs://abcd'
}
}
}
notifier = inventory_summary.InventorySummary(mock_service_config,
'abcd')
notifier._get_output_filename = mock.MagicMock(
return_value=fake_output_name)
gcs_path = '{}/{}'.format('gs://abcd', fake_output_name)
mock_tmp_csv = mock.MagicMock()
mock_tempfile.return_value = mock_tmp_csv
mock_tmp_csv.name = fake_tmpname
mock_tmp_csv.write = mock.MagicMock()
notifier._upload_to_gcs([{}])
mock_tmp_csv.write.assert_called()
mock_storage.return_value.put_text_file.assert_called_once_with(
fake_tmpname, gcs_path)
def test_inventory_summary_no_summary_data(self, mock_logger):
mock_service_config = mock.MagicMock()
mock_service_config.get_notifier_config.return_value = {
'inventory': {
'gcs_summary': {
'enabled': True,
'data_format': 'csv',
'gcs_path': 'gs://blah'
},
'email_summary': {
'enabled': True,
'sendgrid_api_key': 'blah',
'sender': 'blah',
'recipient': 'blah'
}
}
}
notifier = inventory_summary.InventorySummary(mock_service_config,
'abcd')
notifier._get_summary_data = mock.MagicMock()
notifier._get_summary_data.side_effect = util_errors.NoDataError
notifier.run()
self.assertTrue(mock_logger.exception.called)
self.assertTrue('no summary data is found' in
mock_logger.exception.call_args[0][0])
def test_get_gsuite_dwd_status_disabled(self):
mock_inventory_summary_data = [{
'resource_type': 'bucket',
'count': 2
}, {
'resource_type': 'dataset',
'count': 4
}, {
'resource_type': 'folder',
'count': 1
}, {
'resource_type': 'object',
'count': 1
}, {
'resource_type': 'organization',
'count': 1
}, {
'resource_type': 'project',
'count': 2
}]
mock_service_config = mock.MagicMock()
inventory_summary_object = (inventory_summary.InventorySummary(
mock_service_config, 'abcd'))
expected_get_gsuite_dwd_status_output = 'disabled'
test_get_gsuite_dwd_status_output = (
inventory_summary_object._get_gsuite_dwd_status(
mock_inventory_summary_data))
self.assertEqual(expected_get_gsuite_dwd_status_output,
test_get_gsuite_dwd_status_output)
def test_upload_to_gcs_with_invalid_data_format(self, mock_write_csv,
mock_json_stringify,
mock_storage):
"""Test run() with json file format."""
mock_service_config = mock.MagicMock()
mock_service_config.get_notifier_config.return_value = {
'inventory': {
'gcs_summary': {
'enabled': True,
'data_format': 'blah',
'gcs_path': 'gs://abcd'
}
}
}
notifier = inventory_summary.InventorySummary(mock_service_config,
'abcd')
notifier._get_output_filename = mock.MagicMock()
with self.assertRaises(base_notification.InvalidDataFormatError):
notifier._upload_to_gcs([{}])
self.assertFalse(notifier._get_output_filename.called)
self.assertFalse(mock_write_csv.called)
self.assertFalse(mock_json_stringify.called)
def test_get_gsuite_dwd_status_enabled(self):
# arrange
mock_inventory_summary_data = [
{'resource_type': 'bucket', 'count': 2},
{'resource_type': 'dataset', 'count': 4},
{'resource_type': 'folder', 'count': 1},
{'resource_type': 'gsuite_group', 'count': 4},
{'resource_type': 'gsuite_user', 'count': 2},
{'resource_type': 'object', 'count': 1},
{'resource_type': 'organization', 'count': 1},
{'resource_type': 'project', 'count': 2}
]
mock_service_config = mock.MagicMock()
# act
notifier = inventory_summary.InventorySummary(mock_service_config,
'abcd',
mock.MagicMock())
test_get_gsuite_dwd_status_output = (
notifier._get_gsuite_dwd_status(
mock_inventory_summary_data))
# assert
self.assertEqual('enabled',
test_get_gsuite_dwd_status_output)
def test_run_with_json(self, mock_write_csv, mock_json_stringify,
mock_storage):
"""Test run() with json file format."""
mock_json_stringify.return_value = 'test123'
fake_tmpname = 'tmp_name'
fake_output_name = 'abc'
mock_service_config = mock.MagicMock()
mock_service_config.get_notifier_config.return_value = {
'inventory': {
'gcs_summary': {
'enabled': True,
'data_format': 'json',
'gcs_path': 'gs://abcd'
}
}
}
notifier = inventory_summary.InventorySummary(mock_service_config,
'abcd')
notifier._get_output_filename = mock.MagicMock(
return_value=fake_output_name)
notifier._upload_to_gcs([{}])
self.assertTrue(notifier._get_output_filename.called)
self.assertEqual(string_formats.INVENTORY_SUMMARY_JSON_FMT,
notifier._get_output_filename.call_args[0][0])
self.assertFalse(mock_write_csv.called)
self.assertTrue(mock_json_stringify.called)
def test_no_inventory_in_config(self, mock_logger):
mock_service_config = mock.MagicMock()
mock_service_config.get_notifier_config.return_value = dict()
notifier = inventory_summary.InventorySummary(mock_service_config,
'abcd')
notifier.run()
self.assertTrue(mock_logger.info.called)
self.assertEqual('No inventory configuration for notifier.',
mock_logger.info.call_args[0][0])
def test_no_inventory_summary_in_config(self, mock_logger):
mock_service_config = mock.MagicMock()
mock_service_config.get_notifier_config.return_value = {
'inventory': {'blah': 'blah blah'}}
notifier = inventory_summary.InventorySummary(mock_service_config,
'abcd')
notifier.run()
self.assertTrue(
('unable to get inventory summary configuration'
in mock_logger.exception.call_args[0][0]))
def test_inventory_summary_not_enabled_in_config(self, mock_logger):
mock_service_config = mock.MagicMock()
mock_service_config.get_notifier_config.return_value = {
'inventory': {'gcs_summary': {'enabled': False},
'email_summary': {'enabled': False}}}
notifier = inventory_summary.InventorySummary(mock_service_config,
'abcd')
notifier.run()
self.assertTrue(mock_logger.info.called)
self.assertEquals(
'All inventory summaries are disabled.',
mock_logger.info.call_args[0][0])
def test_inventory_summary_invalid_gcs_path(self, mock_logger, mock_write_csv,
mock_json_stringify, mock_storage):
mock_service_config = mock.MagicMock()
mock_service_config.get_notifier_config.return_value = {
'inventory': {'gcs_summary': {'enabled': True,
'data_format': 'blah',
'gcs_path': 'blah'}}}
notifier = inventory_summary.InventorySummary(mock_service_config,
'abcd')
notifier._get_output_filename = mock.MagicMock()
notifier._upload_to_gcs([{}])
self.assertTrue(mock_logger.error.called)
self.assertTrue('Invalid GCS path' in mock_logger.error.call_args[0][0])
def test_inventory_summary_invalid_gcs_path(self, mock_logger):
# arrange
mock_service_config = mock.MagicMock()
mock_service_config.get_notifier_config.return_value = {
'inventory': {'gcs_summary': {'enabled': True,
'data_format': 'blah',
'gcs_path': 'blah'}}}
# act
notifier = inventory_summary.InventorySummary(mock_service_config,
'abcd',
mock.MagicMock())
notifier._get_output_filename = mock.MagicMock()
notifier._upload_to_gcs([{}])
# assert
self.assertTrue(mock_logger.error.called)
self.assertTrue('Invalid GCS path' in mock_logger.error.call_args[0][0])
def test_get_output_filename(self, mock_date_time):
"""Test_get_output_filename()."""
mock_date_time.get_utc_now_datetime = mock.MagicMock()
mock_date_time.get_utc_now_datetime.return_value = self.fake_utcnow
mock_service_config = mock.MagicMock()
mock_service_config.get_notifier_config.return_value = {}
expected_timestamp = self.fake_utcnow.strftime(
string_formats.TIMESTAMP_TIMEZONE_FILES)
notifier = inventory_summary.InventorySummary(mock_service_config,
'abcd')
actual_filename = notifier._get_output_filename(
string_formats.INVENTORY_SUMMARY_CSV_FMT)
expected_filename = string_formats.INVENTORY_SUMMARY_CSV_FMT.format(
notifier.inventory_index_id, expected_timestamp)
self.assertEqual(expected_filename, actual_filename)
def test_gcs_path_is_not_set_in_config(self, mock_logger, mock_write_csv,
mock_json_stringify, mock_storage):
"""Test run() with json file format."""
mock_service_config = mock.MagicMock()
mock_service_config.get_notifier_config.return_value = {
'inventory': {'gcs_summary': {'enabled': True,
'data_format': 'blah'}}}
notifier = inventory_summary.InventorySummary(mock_service_config,
'abcd')
notifier._get_output_filename = mock.MagicMock()
notifier._upload_to_gcs([{}])
self.assertTrue(mock_logger.error.called)
self.assertEquals(
'gcs_path not set for inventory summary notifier.',
mock_logger.error.call_args[0][0])
def test_inventory_summary_can_run_successfully(self):
mock_inventory_index = mock.MagicMock()
mock_inventory_index.get_summary.return_value = {
'bucket': 2,
'dataset': 4,
'folder': 1,
'object': 1,
'organization': 1,
'project': 2
}
mock_inventory_index.get_details.return_value = {
'dataset - HIDDEN': 2,
'dataset - SHOWN': 2,
'project - ACTIVE': 1,
'project - DELETE PENDING': 1
}
mock_session = mock.MagicMock()
mock_session.query.return_value.get.return_value = mock_inventory_index
mock_service_config = mock.MagicMock()
mock_service_config.scoped_session.return_value.__enter__.return_value \
= mock_session
mock_service_config.get_notifier_config.return_value = {
'inventory': {
'gcs_summary': {
'enabled': True,
'data_format': 'csv',
'gcs_path': 'gs://blah'
},
'email_summary': {
'enabled': True,
'sendgrid_api_key': 'blah',
'sender': 'blah',
'recipient': 'blah'
}
}
}
notifier = inventory_summary.InventorySummary(mock_service_config,
'abcd')
mock_get_root_resources = mock.MagicMock()
mock_get_root_resources.return_value = ['111', '222', '333']
notifier._get_root_resources = mock_get_root_resources
notifier._upload_to_gcs = mock.MagicMock()
notifier._send_email = mock.MagicMock()
notifier.run()
expected_summary_data_upload_to_gcs = [{
'count': 2,
'resource_type': 'bucket'
}, {
'count': 4,
'resource_type': 'dataset'
}, {
'count':
2,
'resource_type':
'dataset - HIDDEN'
}, {
'count':
2,
'resource_type':
'dataset - SHOWN'
}, {
'count': 1,
'resource_type': 'folder'
}, {
'count': 1,
'resource_type': 'object'
}, {
'count': 1,
'resource_type': 'organization'
}, {
'count': 2,
'resource_type': 'project'
}, {
'count':
1,
'resource_type':
'project - ACTIVE'
}, {
'count':
1,
'resource_type':
'project - DELETE PENDING'
}]
expected_summary_data_send_email = ([{
'count': 2,
'resource_type': 'bucket'
}, {
'count': 4,
'resource_type': 'dataset'
}, {
'count': 1,
'resource_type': 'folder'
}, {
'count': 1,
'resource_type': 'object'
}, {
'count': 1,
'resource_type': 'organization'
}, {
'count': 2,
'resource_type': 'project'
}], [{
'count': 2,
'resource_type': 'dataset - HIDDEN'
}, {
'count': 2,
'resource_type': 'dataset - SHOWN'
}, {
'count': 1,
'resource_type': 'project - ACTIVE'
}, {
'count': 1,
'resource_type': 'project - DELETE PENDING'
}], ['111', '222', '333'])
self.assertEqual(1, notifier._upload_to_gcs.call_count)
self.assertEqual(expected_summary_data_upload_to_gcs,
notifier._upload_to_gcs.call_args[0][0])
self.assertEqual(1, notifier._send_email.call_count)
self.assertEqual(expected_summary_data_send_email,
notifier._send_email.call_args[0])
