def test_head_method(self, mock_session_send): def resp_check(_): 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(context={}) args, kwargs = mock_session_send.call_args received_request = args[0] prep_request = requests.Request('HEAD', 'https://www.httpbin.org', {}).prepare() self.assertEqual(prep_request.url, received_request.url) self.assertTrue(prep_request.method, received_request.method)
def test_poke_continues_for_http_500_with_extra_options_check_response_false( self, mock_session_send): def resp_check(_): return False response = requests.Response() response.status_code = 500 response.reason = 'Internal Server Error' response._content = b'Internal Server Error' mock_session_send.return_value = response task = HttpSensor( dag=self.dag, task_id='http_sensor_poke_for_code_500', http_conn_id='http_default', endpoint='', request_params={}, method='HEAD', response_check=resp_check, extra_options={'check_response': False}, timeout=5, poke_interval=1, ) with self.assertRaises(AirflowSensorTimeout): task.execute(context={})
def test_logging_head_error_request(self, mock_session_send): def resp_check(_): return True response = requests.Response() response.status_code = 404 response.reason = 'Not Found' response._content = b'This endpoint doesnt exist' 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) calls = [ mock.call('HTTP error: %s', 'Not Found'), mock.call('This endpoint doesnt exist'), mock.call('HTTP error: %s', 'Not Found'), mock.call('This endpoint doesnt exist'), mock.call('HTTP error: %s', 'Not Found'), mock.call('This endpoint doesnt exist'), mock.call('HTTP error: %s', 'Not Found'), mock.call('This endpoint doesnt exist'), mock.call('HTTP error: %s', 'Not Found'), mock.call('This endpoint doesnt exist'), mock.call('HTTP error: %s', 'Not Found'), mock.call('This endpoint doesnt exist'), ] mock_errors.assert_has_calls(calls)
def test_poke_context(self, mock_session_send): response = requests.Response() response.status_code = 200 mock_session_send.return_value = response def resp_check(_, execution_date): if 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, timeout=5, poke_interval=1, dag=self.dag) task_instance = TaskInstance(task=task, execution_date=DEFAULT_DATE) task.execute(task_instance.get_template_context())
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(_): 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.assertRaisesRegex(AirflowException, 'AirflowException raised here!'): task.execute(context={})