ashb commented on a change in pull request #4111: [AIRFLOW-3266] Add AWS Athena 
Operator and hook
URL: https://github.com/apache/incubator-airflow/pull/4111#discussion_r230394621
 
 

 ##########
 File path: tests/contrib/operators/test_aws_athena_operator.py
 ##########
 @@ -48,42 +49,62 @@
 }
 
 
+class Iterator(object):
+
+    def __init__(self, tuple):
+        self.tuple = tuple
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        return self.tuple[randrange(len(self.tuple))]
+
+
 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')
+                                        
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=Iterator(("RUNNING", "SUCCESS",)))
 
 Review comment:
   Random-ness in tests seem like a bad plan at first sight.
   
   Can you explain your reasoning for doing this, rather than say just 
`side_effect=["RUNNING", "SUCCESS"]` which could return running first time and 
success the second time. Right now with randomness we don't know how many times 
it might have to check the query otherwise, or am I missing something?

----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on GitHub and use the
URL above to go to the specific comment.
 
For queries about this service, please contact Infrastructure at:
[email protected]


With regards,
Apache Git Services

Reply via email to