def test_logging_head_error_request(self, mock_session_send): def resp_check(resp): return True response = requests.Response() response.status_code = 404 response.reason = 'Not Found' mock_session_send.return_value = response task = HttpSensor(dag=self.dag, task_id='http_sensor_head_method', http_conn_id='http_default', endpoint='', request_params={}, method='HEAD', response_check=resp_check, timeout=5, poke_interval=1) with mock.patch.object(task.hook.log, 'error') as mock_errors: with self.assertRaises(AirflowSensorTimeout): task.execute(None) self.assertTrue(mock_errors.called) mock_errors.assert_called_with('HTTP error: %s', 'Not Found')
def test_head_method(self, mock_session_send): def resp_check(resp): return True task = HttpSensor( dag=self.dag, task_id='http_sensor_head_method', http_conn_id='http_default', endpoint='', request_params={}, method='HEAD', response_check=resp_check, timeout=5, poke_interval=1) task.execute(None) args, kwargs = mock_session_send.call_args received_request = args[0] prep_request = requests.Request( 'HEAD', 'https://www.google.com', {}).prepare() self.assertEqual(prep_request.url, received_request.url) self.assertTrue(prep_request.method, received_request.method)
def test_logging_head_error_request( self, mock_session_send ): def resp_check(resp): return True response = requests.Response() response.status_code = 404 response.reason = 'Not Found' mock_session_send.return_value = response task = HttpSensor( dag=self.dag, task_id='http_sensor_head_method', http_conn_id='http_default', endpoint='', request_params={}, method='HEAD', response_check=resp_check, timeout=5, poke_interval=1 ) with mock.patch.object(task.hook.log, 'error') as mock_errors: with self.assertRaises(AirflowSensorTimeout): task.execute(None) self.assertTrue(mock_errors.called) mock_errors.assert_called_with('HTTP error: %s', 'Not Found')
def test_poke_exception(self, mock_session_send): """ Exception occurs in poke function should not be ignored. """ response = requests.Response() response.status_code = 200 mock_session_send.return_value = response def resp_check(resp): raise AirflowException('AirflowException raised here!') task = HttpSensor(task_id='http_sensor_poke_exception', http_conn_id='http_default', endpoint='', request_params={}, response_check=resp_check, timeout=5, poke_interval=1) with self.assertRaisesRegexp(AirflowException, 'AirflowException raised here!'): task.execute(None)
def test_poke_exception(self, mock_session_send): """ Exception occurs in poke function should not be ignored. """ response = requests.Response() response.status_code = 200 mock_session_send.return_value = response def resp_check(resp): raise AirflowException('AirflowException raised here!') task = HttpSensor( task_id='http_sensor_poke_exception', http_conn_id='http_default', endpoint='', request_params={}, response_check=resp_check, timeout=5, poke_interval=1) with self.assertRaisesRegexp(AirflowException, 'AirflowException raised here!'): task.execute(None)
def test_poke_context(self, mock_session_send): """ test provide_context """ response = requests.Response() response.status_code = 200 mock_session_send.return_value = response def resp_check(resp, **context): if context: if "execution_date" in context: if context["execution_date"] == DEFAULT_DATE: return True raise AirflowException('AirflowException raised here!') task = HttpSensor(task_id='http_sensor_poke_exception', http_conn_id='http_default', endpoint='', request_params={}, response_check=resp_check, provide_context=True, timeout=5, poke_interval=1, dag=self.dag) task_instance = TaskInstance(task=task, execution_date=DEFAULT_DATE) task.execute(task_instance.get_template_context())