def test_run_athena_task(self) -> None:
mock_client = Mock()
mock_client_factory = Mock(get_athena_client=Mock(return_value=mock_client))
mock_task = Mock(
spec=AwsAthenaTask, run=Mock(side_effect=lambda client: task_report() if client == mock_client else None)
)
self.assertEqual(task_report(), AwsTaskRunner(mock_client_factory)._run_task(mock_task))
def test_run_tasks(self) -> None:
succeeding_task_1 = athena_task(description="some task")
succeeding_task_1._run_task = lambda _: {"outcome_1": "success_1"}
succeeding_task_2 = athena_task(description="other task")
succeeding_task_2._run_task = lambda _: {"outcome_2": "success_2"}
failing_task = athena_task(account=account("5678", "wrong account"),
description="boom")
failing_task._run_task = lambda _: _raise(AwsScannerException("oops"))
tasks = [succeeding_task_1, failing_task, succeeding_task_2]
with self.assertLogs("AwsParallelTaskRunner",
level="ERROR") as error_log:
reports = AwsParallelTaskRunner(Mock()).run(tasks)
self.assertEqual(2, len(reports),
"there should only be two task reports")
self.assertIn(
task_report(description="some task",
results={"outcome_1": "success_1"},
partition=None), reports)
self.assertIn(
task_report(description="other task",
results={"outcome_2": "success_2"},
partition=None), reports)
self.assertEqual(
[
"ERROR:AwsParallelTaskRunner:task 'boom' for 'wrong account (5678)' failed with: 'oops'"
],
error_log.output,
)
def test_run(self) -> None:
tasks = [athena_task(description="task_34"), athena_task(description="task_23")]
report = [task_report(description="task_34"), task_report(description="task_23")]
task_runner = AwsTaskRunner(Mock())
mock_run_tasks = Mock(return_value=report)
with patch.object(task_runner, "_run_tasks", mock_run_tasks):
self.assertEqual(report, task_runner.run(tasks))
mock_run_tasks.assert_called_once_with(tasks)
def _assert_task_run(self, task_type, task_args, query, query_results,
results):
for task_index, query_result in enumerate(query_results):
athena = Mock(run_query=Mock(return_value=query_result))
task = self.__build_task_under_test(task_type=task_type,
task_args=task_args)
self.assertEqual(
task_report(description=task._description,
results=results[task_index]), task.run(athena))
athena.run_query.assert_called_once_with(database="some_db",
query=query)
class TestAwsAthenaCleanerTask(AwsScannerTestCase):
database_mappings = {
"db_1": ["table_1", "table_2", "table_3"],
"some_prefix_db_2": ["table_1", "table_2"],
"db_3": ["table_1"],
"some_prefix_db_4": ["table_1", "table_2", "table_3"],
"some_prefix_db_5": [],
}
expected_report = task_report(
account=account("555666777888", "cloudtrail"),
description="clean scanner leftovers",
partition=None,
results={
"dropped_tables": [
"some_prefix_db_2.table_1",
"some_prefix_db_2.table_2",
"some_prefix_db_4.table_1",
"some_prefix_db_4.table_2",
"some_prefix_db_4.table_3",
],
"dropped_databases":
["some_prefix_db_2", "some_prefix_db_4", "some_prefix_db_5"],
},
)
def test_clean_task_databases(self) -> None:
mock_athena = Mock(
list_databases=Mock(
return_value=list(self.database_mappings.keys())),
list_tables=Mock(
side_effect=lambda db: self.database_mappings.get(db)),
)
self.assertEqual(self.expected_report,
AwsAthenaCleanerTask().run(mock_athena))
mock_athena.assert_has_calls([
call.list_databases(),
call.list_tables("some_prefix_db_2"),
call.drop_table("some_prefix_db_2", "table_1"),
call.drop_table("some_prefix_db_2", "table_2"),
call.list_tables("some_prefix_db_4"),
call.drop_table("some_prefix_db_4", "table_1"),
call.drop_table("some_prefix_db_4", "table_2"),
call.drop_table("some_prefix_db_4", "table_3"),
call.list_tables("some_prefix_db_5"),
call.drop_database("some_prefix_db_2"),
call.drop_database("some_prefix_db_4"),
call.drop_database("some_prefix_db_5"),
])
def test_run(self) -> None:
task_class = "src.tasks.aws_cloudtrail_task.AwsCloudTrailTask"
results = {"key": "val"}
mock_athena = Mock()
mocks = Mock(setup=Mock(),
task=Mock(return_value=results),
teardown=Mock())
with patch(f"{task_class}._setup", mocks.setup):
with patch(f"{task_class}._run_task", mocks.task):
with patch(f"{task_class}._teardown", mocks.teardown):
self.assertEqual(task_report(),
cloudtrail_task().run(mock_athena))
mocks.assert_has_calls([
call.setup(mock_athena),
call.task(mock_athena),
call.teardown(mock_athena)
])
def test_run(self) -> None:
with patch("src.tasks.aws_task.AwsTask._run_task",
return_value={"key": "val"}):
self.assertEqual(task_report(partition=None),
aws_task().run(Mock()))
def test_run_s3_task(self) -> None:
client = Mock()
client_factory = Mock(get_s3_client=Mock(side_effect=lambda acc: client if acc == account() else None))
task = s3_task()
task.run = Mock(side_effect=lambda c: task_report() if c == client else None) # type: ignore
self.assertEqual(task_report(), AwsTaskRunner(client_factory)._run_task(task))
def test_serialize_complex_object_with_optional_properties(self) -> None:
self.assertEqual(
('{"account": {"identifier": "account_id", "name": "account_name"}, "description": "task", '
'"results": {"key": "val"}}'),
to_json(task_report(partition=None)),
)
def build_test_report(description):
return [task_report(description=description, results={"key": "val"})]
