def get_hook(self):
"""Create and return an AWSAthenaHook."""
return AWSAthenaHook(self.aws_conn_id, sleep_time=self.sleep_time)
def setUp(self):
self.athena = AWSAthenaHook(sleep_time=0)
class TestAWSAthenaHook(unittest.TestCase):
def setUp(self):
self.athena = AWSAthenaHook(sleep_time=0)
def test_init(self):
self.assertEqual(self.athena.aws_conn_id, 'aws_default')
self.assertEqual(self.athena.sleep_time, 0)
@mock.patch.object(AWSAthenaHook, 'get_conn')
def test_hook_run_query_without_token(self, mock_conn):
mock_conn.return_value.start_query_execution.return_value = MOCK_QUERY_EXECUTION
result = self.athena.run_query(
query=MOCK_DATA['query'],
query_context=mock_query_context,
result_configuration=mock_result_configuration,
)
expected_call_params = {
'QueryString': MOCK_DATA['query'],
'QueryExecutionContext': mock_query_context,
'ResultConfiguration': mock_result_configuration,
'WorkGroup': MOCK_DATA['workgroup'],
}
mock_conn.return_value.start_query_execution.assert_called_with(**expected_call_params)
self.assertEqual(result, MOCK_DATA['query_execution_id'])
@mock.patch.object(AWSAthenaHook, 'get_conn')
def test_hook_run_query_with_token(self, mock_conn):
mock_conn.return_value.start_query_execution.return_value = MOCK_QUERY_EXECUTION
result = self.athena.run_query(
query=MOCK_DATA['query'],
query_context=mock_query_context,
result_configuration=mock_result_configuration,
client_request_token=MOCK_DATA['client_request_token'],
)
expected_call_params = {
'QueryString': MOCK_DATA['query'],
'QueryExecutionContext': mock_query_context,
'ResultConfiguration': mock_result_configuration,
'ClientRequestToken': MOCK_DATA['client_request_token'],
'WorkGroup': MOCK_DATA['workgroup'],
}
mock_conn.return_value.start_query_execution.assert_called_with(**expected_call_params)
self.assertEqual(result, MOCK_DATA['query_execution_id'])
@mock.patch.object(AWSAthenaHook, 'get_conn')
def test_hook_get_query_results_with_non_succeeded_query(self, mock_conn):
mock_conn.return_value.get_query_execution.return_value = MOCK_RUNNING_QUERY_EXECUTION
result = self.athena.get_query_results(query_execution_id=MOCK_DATA['query_execution_id'])
self.assertIsNone(result)
@mock.patch.object(AWSAthenaHook, 'get_conn')
def test_hook_get_query_results_with_default_params(self, mock_conn):
mock_conn.return_value.get_query_execution.return_value = MOCK_SUCCEEDED_QUERY_EXECUTION
self.athena.get_query_results(query_execution_id=MOCK_DATA['query_execution_id'])
expected_call_params = {'QueryExecutionId': MOCK_DATA['query_execution_id'], 'MaxResults': 1000}
mock_conn.return_value.get_query_results.assert_called_with(**expected_call_params)
@mock.patch.object(AWSAthenaHook, 'get_conn')
def test_hook_get_query_results_with_next_token(self, mock_conn):
mock_conn.return_value.get_query_execution.return_value = MOCK_SUCCEEDED_QUERY_EXECUTION
self.athena.get_query_results(
query_execution_id=MOCK_DATA['query_execution_id'], next_token_id=MOCK_DATA['next_token_id']
)
expected_call_params = {
'QueryExecutionId': MOCK_DATA['query_execution_id'],
'NextToken': MOCK_DATA['next_token_id'],
'MaxResults': 1000,
}
mock_conn.return_value.get_query_results.assert_called_with(**expected_call_params)
@mock.patch.object(AWSAthenaHook, 'get_conn')
def test_hook_get_paginator_with_non_succeeded_query(self, mock_conn):
mock_conn.return_value.get_query_execution.return_value = MOCK_RUNNING_QUERY_EXECUTION
result = self.athena.get_query_results_paginator(query_execution_id=MOCK_DATA['query_execution_id'])
self.assertIsNone(result)
@mock.patch.object(AWSAthenaHook, 'get_conn')
def test_hook_get_paginator_with_default_params(self, mock_conn):
mock_conn.return_value.get_query_execution.return_value = MOCK_SUCCEEDED_QUERY_EXECUTION
self.athena.get_query_results_paginator(query_execution_id=MOCK_DATA['query_execution_id'])
expected_call_params = {
'QueryExecutionId': MOCK_DATA['query_execution_id'],
'PaginationConfig': {'MaxItems': None, 'PageSize': None, 'StartingToken': None},
}
mock_conn.return_value.get_paginator.return_value.paginate.assert_called_with(**expected_call_params)
@mock.patch.object(AWSAthenaHook, 'get_conn')
def test_hook_get_paginator_with_pagination_config(self, mock_conn):
mock_conn.return_value.get_query_execution.return_value = MOCK_SUCCEEDED_QUERY_EXECUTION
self.athena.get_query_results_paginator(
query_execution_id=MOCK_DATA['query_execution_id'],
max_items=MOCK_DATA['max_items'],
page_size=MOCK_DATA['max_items'],
starting_token=MOCK_DATA['next_token_id'],
)
expected_call_params = {
'QueryExecutionId': MOCK_DATA['query_execution_id'],
'PaginationConfig': {
'MaxItems': MOCK_DATA['max_items'],
'PageSize': MOCK_DATA['max_items'],
'StartingToken': MOCK_DATA['next_token_id'],
},
}
mock_conn.return_value.get_paginator.return_value.paginate.assert_called_with(**expected_call_params)
@mock.patch.object(AWSAthenaHook, 'get_conn')
def test_hook_poll_query_when_final(self, mock_conn):
mock_conn.return_value.get_query_execution.return_value = MOCK_SUCCEEDED_QUERY_EXECUTION
result = self.athena.poll_query_status(query_execution_id=MOCK_DATA['query_execution_id'])
mock_conn.return_value.get_query_execution.assert_called_once()
self.assertEqual(result, 'SUCCEEDED')
@mock.patch.object(AWSAthenaHook, 'get_conn')
def test_hook_poll_query_with_timeout(self, mock_conn):
mock_conn.return_value.get_query_execution.return_value = MOCK_RUNNING_QUERY_EXECUTION
result = self.athena.poll_query_status(
query_execution_id=MOCK_DATA['query_execution_id'], max_tries=1
)
mock_conn.return_value.get_query_execution.assert_called_once()
self.assertEqual(result, 'RUNNING')
def get_hook(self):
return AWSAthenaHook(self.aws_conn_id, self.sleep_time)
def hook(self) -> AWSAthenaHook:
"""Create and return an AWSAthenaHook"""
return AWSAthenaHook(self.aws_conn_id, sleep_time=self.sleep_time)
def get_hook(self):
"""Create and return an AWSAthenaHook"""
if not self.hook:
self.hook = AWSAthenaHook(self.aws_conn_id, self.sleep_time)
return self.hook
