Correctly pass sleep time from AWSAthenaOperator down to the hook. (#8845)
Sleep time in AthenaHook was defined as a kwarg-only key. This one change makes tests go from 270s to 0.3s :D
This commit is contained in:
Родитель
8a94d18c04
Коммит
ed3f5131a2
|
@ -81,7 +81,7 @@ class AWSAthenaOperator(BaseOperator):
|
||||||
|
|
||||||
def get_hook(self):
|
def get_hook(self):
|
||||||
"""Create and return an AWSAthenaHook."""
|
"""Create and return an AWSAthenaHook."""
|
||||||
return AWSAthenaHook(self.aws_conn_id, self.sleep_time)
|
return AWSAthenaHook(self.aws_conn_id, sleep_time=self.sleep_time)
|
||||||
|
|
||||||
def execute(self, context):
|
def execute(self, context):
|
||||||
"""
|
"""
|
||||||
|
|
|
@ -62,7 +62,7 @@ class TestAWSAthenaOperator(unittest.TestCase):
|
||||||
self.athena = AWSAthenaOperator(task_id='test_aws_athena_operator', query='SELECT * FROM TEST_TABLE',
|
self.athena = AWSAthenaOperator(task_id='test_aws_athena_operator', query='SELECT * FROM TEST_TABLE',
|
||||||
database='TEST_DATABASE', output_location='s3://test_s3_bucket/',
|
database='TEST_DATABASE', output_location='s3://test_s3_bucket/',
|
||||||
client_request_token='eac427d0-1c6d-4dfb-96aa-2835d3ac6595',
|
client_request_token='eac427d0-1c6d-4dfb-96aa-2835d3ac6595',
|
||||||
sleep_time=1, max_tries=3, dag=self.dag)
|
sleep_time=0, max_tries=3, dag=self.dag)
|
||||||
|
|
||||||
def test_init(self):
|
def test_init(self):
|
||||||
self.assertEqual(self.athena.task_id, MOCK_DATA['task_id'])
|
self.assertEqual(self.athena.task_id, MOCK_DATA['task_id'])
|
||||||
|
@ -70,7 +70,10 @@ class TestAWSAthenaOperator(unittest.TestCase):
|
||||||
self.assertEqual(self.athena.database, MOCK_DATA['database'])
|
self.assertEqual(self.athena.database, MOCK_DATA['database'])
|
||||||
self.assertEqual(self.athena.aws_conn_id, 'aws_default')
|
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.client_request_token, MOCK_DATA['client_request_token'])
|
||||||
self.assertEqual(self.athena.sleep_time, 1)
|
self.assertEqual(self.athena.sleep_time, 0)
|
||||||
|
|
||||||
|
hook = self.athena.get_hook()
|
||||||
|
self.assertEqual(hook.sleep_time, 0)
|
||||||
|
|
||||||
@mock.patch.object(AWSAthenaHook, 'check_query_status', side_effect=("SUCCESS",))
|
@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, 'run_query', return_value=ATHENA_QUERY_ID)
|
||||||
|
|
Загрузка…
Ссылка в новой задаче