def test_get_results_with_headers(self):
dag = mock.MagicMock()
dag.dag_id = DAG_ID
hook = QuboleHook(task_id=TASK_ID, command_type='prestocmd', dag=dag)
task = mock.MagicMock()
task.xcom_pull.return_value = 'test_command_id'
with mock.patch('qds_sdk.resource.Resource.find',
return_value=PrestoCommand):
results = open(hook.get_results(ti=task,
include_headers=True)).read()
assert results == RESULTS_WITH_HEADER
def get_hook(self) -> QuboleHook:
"""Reinitialising the hook, as some template fields might have changed"""
return QuboleHook(**self.kwargs)