[GitHub] [airflow] eladkal commented on a change in pull request #15683: Feature qubole hook support headers

[GitBox <gi...@apache.org>]

eladkal commented on a change in pull request #15683:
URL: https://github.com/apache/airflow/pull/15683#discussion_r628420689



##########
File path: tests/providers/qubole/operators/test_qubole.py
##########
@@ -180,3 +192,16 @@ def test_parameter_pool_passed(self):
         test_pool = 'test_pool'
         op = QuboleOperator(task_id=TASK_ID, pool=test_pool)
         assert op.pool == test_pool
+
+    @mock.patch('qds_sdk.commands.Command.get_results', new=get_result_mock)
+    def test_get_results_with_headers(self):
+        dag = DAG(DAG_ID, start_date=DEFAULT_DATE)
+
+        task = QuboleOperator(task_id=TASK_ID, command_type='prestocmd', dag=dag)
+
+        with mock.patch.object(task, 'xcom_pull', return_value='test_command_id'):
+            with mock.patch('qds_sdk.resource.Resource.find', return_value=PrestoCommand):
+                results = open(task.get_results(ti=task, include_headers=True), 'r').read()
+                assert results == RESULTS_WITH_HEADER
+                results = open(task.get_results(ti=task, include_headers=False), 'r').read()
+                assert results == RESULTS_WITH_NO_HEADER

Review comment:
       The two asserts actually test two different test cases so lets you seperate it to 2 tests




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org