diff --git a/airflow/providers/tableau/operators/tableau.py b/airflow/providers/tableau/operators/tableau.py index f1547b0fac26e..1145f008bb912 100644 --- a/airflow/providers/tableau/operators/tableau.py +++ b/airflow/providers/tableau/operators/tableau.py @@ -116,16 +116,16 @@ def execute(self, context: dict) -> str: response = method(resource_id) - job_id = response.id - - if self.method == 'refresh': - if self.blocking_refresh: - if not tableau_hook.wait_for_state( - job_id=job_id, - check_interval=self.check_interval, - target_state=TableauJobFinishCode.SUCCESS, - ): - raise TableauJobFailedException(f'The Tableau Refresh {self.resource} Job failed!') + job_id = response.id + + if self.method == 'refresh': + if self.blocking_refresh: + if not tableau_hook.wait_for_state( + job_id=job_id, + check_interval=self.check_interval, + target_state=TableauJobFinishCode.SUCCESS, + ): + raise TableauJobFailedException(f'The Tableau Refresh {self.resource} Job failed!') return job_id diff --git a/tests/providers/tableau/operators/test_tableau.py b/tests/providers/tableau/operators/test_tableau.py index 4d8e20f5e0228..b06244fb61dc0 100644 --- a/tests/providers/tableau/operators/test_tableau.py +++ b/tests/providers/tableau/operators/test_tableau.py @@ -76,8 +76,26 @@ def test_execute_workbooks_blocking(self, mock_tableau_hook): """ Test execute workbooks blocking """ + mock_signed_in = [False] + + def mock_hook_enter(): + mock_signed_in[0] = True + return mock_tableau_hook + + def mock_hook_exit(exc_type, exc_val, exc_tb): + mock_signed_in[0] = False + + def mock_wait_for_state(job_id, target_state, check_interval): + if not mock_signed_in[0]: + raise Exception('Not signed in') + + return True + + mock_tableau_hook.return_value.__enter__ = Mock(side_effect=mock_hook_enter) + mock_tableau_hook.return_value.__exit__ = Mock(side_effect=mock_hook_exit) + mock_tableau_hook.wait_for_state = Mock(side_effect=mock_wait_for_state) + mock_tableau_hook.get_all = Mock(return_value=self.mocked_workbooks) - mock_tableau_hook.return_value.__enter__ = Mock(return_value=mock_tableau_hook) mock_tableau_hook.server.jobs.get_by_id = Mock( return_value=Mock(finish_code=TableauJobFinishCode.SUCCESS.value) ) @@ -123,8 +141,26 @@ def test_execute_datasources_blocking(self, mock_tableau_hook): """ Test execute datasources blocking """ + mock_signed_in = [False] + + def mock_hook_enter(): + mock_signed_in[0] = True + return mock_tableau_hook + + def mock_hook_exit(exc_type, exc_val, exc_tb): + mock_signed_in[0] = False + + def mock_wait_for_state(job_id, target_state, check_interval): + if not mock_signed_in[0]: + raise Exception('Not signed in') + + return True + + mock_tableau_hook.return_value.__enter__ = Mock(side_effect=mock_hook_enter) + mock_tableau_hook.return_value.__exit__ = Mock(side_effect=mock_hook_exit) + mock_tableau_hook.wait_for_state = Mock(side_effect=mock_wait_for_state) + mock_tableau_hook.get_all = Mock(return_value=self.mock_datasources) - mock_tableau_hook.return_value.__enter__ = Mock(return_value=mock_tableau_hook) operator = TableauOperator(find='ds_2', resource='datasources', **self.kwargs) job_id = operator.execute(context={})