From ed3f5131a27e2ef0422f2495a4532630a6204f82 Mon Sep 17 00:00:00 2001 From: Ash Berlin-Taylor Date: Wed, 13 May 2020 09:37:37 +0100 Subject: [PATCH] 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 --- airflow/providers/amazon/aws/operators/athena.py | 2 +- tests/providers/amazon/aws/operators/test_athena.py | 7 +++++-- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/airflow/providers/amazon/aws/operators/athena.py b/airflow/providers/amazon/aws/operators/athena.py index f2724f3c63..28781aa41a 100644 --- a/airflow/providers/amazon/aws/operators/athena.py +++ b/airflow/providers/amazon/aws/operators/athena.py @@ -81,7 +81,7 @@ class AWSAthenaOperator(BaseOperator): def get_hook(self): """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): """ diff --git a/tests/providers/amazon/aws/operators/test_athena.py b/tests/providers/amazon/aws/operators/test_athena.py index 21672a811f..109c983462 100644 --- a/tests/providers/amazon/aws/operators/test_athena.py +++ b/tests/providers/amazon/aws/operators/test_athena.py @@ -62,7 +62,7 @@ class TestAWSAthenaOperator(unittest.TestCase): 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) + sleep_time=0, max_tries=3, dag=self.dag) def test_init(self): 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.aws_conn_id, 'aws_default') 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, 'run_query', return_value=ATHENA_QUERY_ID)