class TestAWSAthenaOperator(unittest.TestCase): def setUp(self): configuration.load_test_config() self.athena = AWSAthenaOperator(task_id='test_aws_athena_operator', query='SELECT * FROM TEST_TABLE', database='TEST_DATABASE', output_location='s3://test_s3_bucket/', client_request_token='eac427d0-1c6d-4dfb-96aa-2835d3ac6595', sleep_time=1) def test_init(self): self.assertEqual(self.athena.task_id, MOCK_DATA['task_id']) self.assertEqual(self.athena.query, MOCK_DATA['query']) self.assertEqual(self.athena.database, MOCK_DATA['database']) self.assertEqual(self.athena.aws_conn_id, 'aws_default') self.assertEqual(self.athena.client_request_token, MOCK_DATA['client_request_token']) self.assertEqual(self.athena.sleep_time, 1) @mock.patch.object(AWSAthenaHook, 'check_query_status', side_effect=("SUCCESS",)) @mock.patch.object(AWSAthenaHook, 'run_query', return_value='1234') @mock.patch.object(AWSAthenaHook, 'get_conn') def test_hook_run_small_success_query(self, mock_conn, mock_run_query, mock_check_query_status): self.athena.execute(None) mock_run_query.assert_called_once_with(MOCK_DATA['query'], query_context, result_configuration, MOCK_DATA['client_request_token']) self.assertEqual(mock_check_query_status.call_count, 1) @mock.patch.object(AWSAthenaHook, 'check_query_status', side_effect=("RUNNING", "RUNNING", "SUCCESS",)) @mock.patch.object(AWSAthenaHook, 'run_query', return_value='1234') @mock.patch.object(AWSAthenaHook, 'get_conn') def test_hook_run_big_success_query(self, mock_conn, mock_run_query, mock_check_query_status): self.athena.execute(None) mock_run_query.assert_called_once_with(MOCK_DATA['query'], query_context, result_configuration, MOCK_DATA['client_request_token']) self.assertEqual(mock_check_query_status.call_count, 3) @mock.patch.object(AWSAthenaHook, 'check_query_status', side_effect=("RUNNING", "FAILED",)) @mock.patch.object(AWSAthenaHook, 'run_query', return_value='1234') @mock.patch.object(AWSAthenaHook, 'get_conn') def test_hook_run_failure_query(self, mock_conn, mock_run_query, mock_check_query_status): self.athena.execute(None) mock_run_query.assert_called_once_with(MOCK_DATA['query'], query_context, result_configuration, MOCK_DATA['client_request_token']) self.assertEqual(mock_check_query_status.call_count, 2) @mock.patch.object(AWSAthenaHook, 'check_query_status', side_effect=("RUNNING", "RUNNING", "CANCELLED",)) @mock.patch.object(AWSAthenaHook, 'run_query', return_value='1234') @mock.patch.object(AWSAthenaHook, 'get_conn') def test_hook_run_cancelled_query(self, mock_conn, mock_run_query, mock_check_query_status): self.athena.execute(None) mock_run_query.assert_called_once_with(MOCK_DATA['query'], query_context, result_configuration, MOCK_DATA['client_request_token']) self.assertEqual(mock_check_query_status.call_count, 3)
class TestAWSAthenaOperator(unittest.TestCase): def setUp(self): args = { 'owner': 'airflow', 'start_date': DEFAULT_DATE, 'provide_context': True } self.dag = DAG(TEST_DAG_ID + 'test_schedule_dag_once', default_args=args, schedule_interval='@once') self.athena = AWSAthenaOperator( task_id='test_aws_athena_operator', query='SELECT * FROM TEST_TABLE', database='TEST_DATABASE', output_location='s3://test_s3_bucket/', client_request_token='eac427d0-1c6d-4dfb-96aa-2835d3ac6595', sleep_time=1, max_tries=3, dag=self.dag) def test_init(self): self.assertEqual(self.athena.task_id, MOCK_DATA['task_id']) self.assertEqual(self.athena.query, MOCK_DATA['query']) self.assertEqual(self.athena.database, MOCK_DATA['database']) self.assertEqual(self.athena.aws_conn_id, 'aws_default') self.assertEqual(self.athena.client_request_token, MOCK_DATA['client_request_token']) self.assertEqual(self.athena.sleep_time, 1) @mock.patch.object(AWSAthenaHook, 'check_query_status', side_effect=("SUCCESS", )) @mock.patch.object(AWSAthenaHook, 'run_query', return_value=ATHENA_QUERY_ID) @mock.patch.object(AWSAthenaHook, 'get_conn') def test_hook_run_small_success_query(self, mock_conn, mock_run_query, mock_check_query_status): self.athena.execute(None) mock_run_query.assert_called_once_with( MOCK_DATA['query'], query_context, result_configuration, MOCK_DATA['client_request_token'], MOCK_DATA['workgroup']) self.assertEqual(mock_check_query_status.call_count, 1) @mock.patch.object(AWSAthenaHook, 'check_query_status', side_effect=( "RUNNING", "RUNNING", "SUCCESS", )) @mock.patch.object(AWSAthenaHook, 'run_query', return_value=ATHENA_QUERY_ID) @mock.patch.object(AWSAthenaHook, 'get_conn') def test_hook_run_big_success_query(self, mock_conn, mock_run_query, mock_check_query_status): self.athena.execute(None) mock_run_query.assert_called_once_with( MOCK_DATA['query'], query_context, result_configuration, MOCK_DATA['client_request_token'], MOCK_DATA['workgroup']) self.assertEqual(mock_check_query_status.call_count, 3) @mock.patch.object(AWSAthenaHook, 'check_query_status', side_effect=( None, None, )) @mock.patch.object(AWSAthenaHook, 'run_query', return_value=ATHENA_QUERY_ID) @mock.patch.object(AWSAthenaHook, 'get_conn') def test_hook_run_failed_query_with_none(self, mock_conn, mock_run_query, mock_check_query_status): with self.assertRaises(Exception): self.athena.execute(None) mock_run_query.assert_called_once_with( MOCK_DATA['query'], query_context, result_configuration, MOCK_DATA['client_request_token'], MOCK_DATA['workgroup']) self.assertEqual(mock_check_query_status.call_count, 3) @mock.patch.object(AWSAthenaHook, 'get_state_change_reason') @mock.patch.object(AWSAthenaHook, 'check_query_status', side_effect=( "RUNNING", "FAILED", )) @mock.patch.object(AWSAthenaHook, 'run_query', return_value=ATHENA_QUERY_ID) @mock.patch.object(AWSAthenaHook, 'get_conn') def test_hook_run_failure_query(self, mock_conn, mock_run_query, mock_check_query_status, mock_get_state_change_reason): with self.assertRaises(Exception): self.athena.execute(None) mock_run_query.assert_called_once_with( MOCK_DATA['query'], query_context, result_configuration, MOCK_DATA['client_request_token'], MOCK_DATA['workgroup']) self.assertEqual(mock_check_query_status.call_count, 2) self.assertEqual(mock_get_state_change_reason.call_count, 1) @mock.patch.object(AWSAthenaHook, 'check_query_status', side_effect=( "RUNNING", "RUNNING", "CANCELLED", )) @mock.patch.object(AWSAthenaHook, 'run_query', return_value=ATHENA_QUERY_ID) @mock.patch.object(AWSAthenaHook, 'get_conn') def test_hook_run_cancelled_query(self, mock_conn, mock_run_query, mock_check_query_status): with self.assertRaises(Exception): self.athena.execute(None) mock_run_query.assert_called_once_with( MOCK_DATA['query'], query_context, result_configuration, MOCK_DATA['client_request_token'], MOCK_DATA['workgroup']) self.assertEqual(mock_check_query_status.call_count, 3) @mock.patch.object(AWSAthenaHook, 'check_query_status', side_effect=( "RUNNING", "RUNNING", "RUNNING", )) @mock.patch.object(AWSAthenaHook, 'run_query', return_value=ATHENA_QUERY_ID) @mock.patch.object(AWSAthenaHook, 'get_conn') def test_hook_run_failed_query_with_max_tries(self, mock_conn, mock_run_query, mock_check_query_status): with self.assertRaises(Exception): self.athena.execute(None) mock_run_query.assert_called_once_with( MOCK_DATA['query'], query_context, result_configuration, MOCK_DATA['client_request_token'], MOCK_DATA['workgroup']) self.assertEqual(mock_check_query_status.call_count, 3) @mock.patch.object(AWSAthenaHook, 'check_query_status', side_effect=("SUCCESS", )) @mock.patch.object(AWSAthenaHook, 'run_query', return_value=ATHENA_QUERY_ID) @mock.patch.object(AWSAthenaHook, 'get_conn') def test_xcom_push_and_pull(self, mock_conn, mock_run_query, mock_check_query_status): ti = TaskInstance(task=self.athena, execution_date=timezone.utcnow()) ti.run() self.assertEqual(ti.xcom_pull(task_ids='test_aws_athena_operator'), ATHENA_QUERY_ID)
class TestAWSAthenaOperator(unittest.TestCase): def setUp(self): configuration.load_test_config() self.athena = AWSAthenaOperator( task_id='test_aws_athena_operator', query='SELECT * FROM TEST_TABLE', database='TEST_DATABASE', output_location='s3://test_s3_bucket/', client_request_token='eac427d0-1c6d-4dfb-96aa-2835d3ac6595', sleep_time=1) def test_init(self): self.assertEqual(self.athena.task_id, MOCK_DATA['task_id']) self.assertEqual(self.athena.query, MOCK_DATA['query']) self.assertEqual(self.athena.database, MOCK_DATA['database']) self.assertEqual(self.athena.aws_conn_id, 'aws_default') self.assertEqual(self.athena.client_request_token, MOCK_DATA['client_request_token']) self.assertEqual(self.athena.sleep_time, 1) @mock.patch.object(AWSAthenaHook, 'check_query_status', side_effect=("SUCCESS", )) @mock.patch.object(AWSAthenaHook, 'run_query', return_value='1234') @mock.patch.object(AWSAthenaHook, 'get_conn') def test_hook_run_small_success_query(self, mock_conn, mock_run_query, mock_check_query_status): self.athena.execute(None) mock_run_query.assert_called_once_with( MOCK_DATA['query'], query_context, result_configuration, MOCK_DATA['client_request_token']) self.assertEqual(mock_check_query_status.call_count, 1) @mock.patch.object(AWSAthenaHook, 'check_query_status', side_effect=( "RUNNING", "RUNNING", "SUCCESS", )) @mock.patch.object(AWSAthenaHook, 'run_query', return_value='1234') @mock.patch.object(AWSAthenaHook, 'get_conn') def test_hook_run_big_success_query(self, mock_conn, mock_run_query, mock_check_query_status): self.athena.execute(None) mock_run_query.assert_called_once_with( MOCK_DATA['query'], query_context, result_configuration, MOCK_DATA['client_request_token']) self.assertEqual(mock_check_query_status.call_count, 3) @mock.patch.object(AWSAthenaHook, 'check_query_status', side_effect=( "RUNNING", "FAILED", )) @mock.patch.object(AWSAthenaHook, 'run_query', return_value='1234') @mock.patch.object(AWSAthenaHook, 'get_conn') def test_hook_run_failure_query(self, mock_conn, mock_run_query, mock_check_query_status): self.athena.execute(None) mock_run_query.assert_called_once_with( MOCK_DATA['query'], query_context, result_configuration, MOCK_DATA['client_request_token']) self.assertEqual(mock_check_query_status.call_count, 2) @mock.patch.object(AWSAthenaHook, 'check_query_status', side_effect=( "RUNNING", "RUNNING", "CANCELLED", )) @mock.patch.object(AWSAthenaHook, 'run_query', return_value='1234') @mock.patch.object(AWSAthenaHook, 'get_conn') def test_hook_run_cancelled_query(self, mock_conn, mock_run_query, mock_check_query_status): self.athena.execute(None) mock_run_query.assert_called_once_with( MOCK_DATA['query'], query_context, result_configuration, MOCK_DATA['client_request_token']) self.assertEqual(mock_check_query_status.call_count, 3)