ColtenOuO commented on code in PR #72133: URL: https://github.com/apache/airflow/pull/72133#discussion_r3872643774
########## providers/standard/tests/unit/standard/sensors/test_websocket.py: ########## @@ -0,0 +1,119 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +from unittest import mock + +import pytest +from websockets.sync.client import ClientConnection + +from airflow.models.dag import DAG +from airflow.providers.common.compat.sdk import AirflowSensorTimeout, TaskDeferred +from airflow.providers.standard.sensors.websocket import WebSocketSensor +from airflow.providers.standard.triggers.websocket import WebSocketTrigger + +from tests_common.test_utils.version_compat import timezone + +URL = "ws://example.com/socket" +DEFAULT_DATE = timezone.datetime(2015, 1, 1) + + +class TestWebSocketSensor: + @classmethod + def setup_class(cls): + args = {"owner": "airflow", "start_date": DEFAULT_DATE} + cls.dag = DAG("test_websocket_sensor", schedule=None, default_args=args) + + @mock.patch("airflow.providers.standard.sensors.websocket.connect", autospec=True) + def test_poke_returns_true_on_message(self, mock_connect): + mock_websocket = mock.MagicMock(spec=ClientConnection) + mock_websocket.recv.return_value = "pong" + mock_connect.return_value.__enter__.return_value = mock_websocket + + sensor = WebSocketSensor(task_id="poke_true", url=URL, message_to_send="ping", dag=self.dag) + assert sensor.poke(context={}) is True + mock_websocket.send.assert_called_once_with("ping") + + @mock.patch("airflow.providers.standard.sensors.websocket.connect", autospec=True) + def test_poke_returns_false_on_timeout(self, mock_connect): + mock_websocket = mock.MagicMock(spec=ClientConnection) + mock_websocket.recv.side_effect = TimeoutError() + mock_connect.return_value.__enter__.return_value = mock_websocket + + sensor = WebSocketSensor(task_id="poke_false", url=URL, dag=self.dag) + assert sensor.poke(context={}) is False + + @mock.patch("airflow.providers.standard.sensors.websocket.connect", autospec=True) + def test_poke_waits_for_the_overall_timeout_not_poke_interval(self, mock_connect): + """recv() must be bounded by the sensor's overall timeout, not poke_interval — + otherwise a single poke could block past the sensor's declared timeout before + that timeout is ever checked.""" + mock_websocket = mock.MagicMock(spec=ClientConnection) + mock_websocket.recv.return_value = "pong" + mock_connect.return_value.__enter__.return_value = mock_websocket + + sensor = WebSocketSensor( + task_id="poke_timeout_arg", url=URL, timeout=45, poke_interval=5, dag=self.dag + ) + assert sensor.poke(context={}) is True + mock_websocket.recv.assert_called_once_with(timeout=45) + + def test_reschedule_mode_not_allowed(self): + with pytest.raises(ValueError, match="Cannot set mode to 'reschedule'. Only 'poke' is acceptable"): + WebSocketSensor(task_id="reschedule", url=URL, mode="reschedule", dag=self.dag) + + def test_task_defer_does_not_poke_first(self): + """The deferrable path must defer immediately: poke() consumes the connection, + so polling before deferring would send message_to_send and lose the reply the + trigger is supposed to wait for.""" + sensor = WebSocketSensor(task_id="defer", url=URL, deferrable=True, dag=self.dag) + + with mock.patch.object(WebSocketSensor, "poke") as mock_poke: Review Comment: The WebSocket mocks above now use spec/autospec, but these two patch.object(WebSocketSensor, "poke") calls still create unspecced mocks. Per Airflow’s testing guidelines, please add autospec=True to both patches so an invalid method signature cannot pass unnoticed. ########## providers/standard/src/airflow/providers/standard/sensors/websocket.py: ########## @@ -0,0 +1,105 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import datetime +from collections.abc import Sequence +from typing import TYPE_CHECKING, Any + +from websockets.sync.client import connect + +from airflow.providers.common.compat.sdk import BaseSensorOperator, conf, poke_mode_only +from airflow.providers.standard.triggers.websocket import WebSocketTrigger + +if TYPE_CHECKING: + from airflow.sdk import Context + + +@poke_mode_only +class WebSocketSensor(BaseSensorOperator): + """ + Waits for a message on a WebSocket connection. + + WebSocket messages are consumptive: once read, a message cannot be read again, and + reconnecting can duplicate ``message_to_send`` against the remote server. The + non-deferrable path therefore opens exactly one connection and blocks on it for up to + ``timeout`` seconds instead of reconnecting every ``poke_interval`` (which has no + effect on this sensor), and this sensor is marked poke-mode-only since a rescheduled + invocation would need a new connection anyway. + + :param url: The ``ws://`` or ``wss://`` URL of the WebSocket server to connect to. + :param header: Optional headers sent when opening the connection. + :param message_to_send: Optional message sent right after the connection is established. + :param deferrable: If waiting for completion, whether to defer the task until done, + default is ``False``. + + .. seealso:: + For more information on how to use this sensor, take a look at the guide: + :ref:`howto/operator:WebSocketSensor` + """ + + template_fields: Sequence[str] = ("url", "header", "message_to_send") + template_fields_renderers = {"header": "json"} + + def __init__( + self, + *, + url: str, + header: dict[str, str] | None = None, + message_to_send: str | bytes | None = None, + deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False), + **kwargs, + ): + super().__init__(**kwargs) + self.url = url + self.header = header + self.message_to_send = message_to_send + self.deferrable = deferrable + + def poke(self, context: Context) -> bool: + self.log.info("Connecting to WebSocket %s", self.url) + with connect(self.url, additional_headers=self.header) as websocket: Review Comment: connect() doesn’t receive an open_timeout, so it uses websockets’ 10-second default. With a sensor timeout=3, the opening handshake may still wait for 10 seconds; if it succeeds, recv() then waits for the full self.timeout again, so the total duration can exceed the declared timeout. A handshake TimeoutError also occurs outside the current try block, so it bypasses the normal sensor-timeout path and soft_fail=True won’t apply. Please create one deadline with time.monotonic(), pass the remaining duration to both connect(open_timeout=...) and recv(timeout=...), and handle timeouts from both stages consistently. ########## providers/standard/tests/unit/standard/sensors/test_websocket.py: ########## @@ -0,0 +1,119 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +from unittest import mock + +import pytest +from websockets.sync.client import ClientConnection + +from airflow.models.dag import DAG +from airflow.providers.common.compat.sdk import AirflowSensorTimeout, TaskDeferred +from airflow.providers.standard.sensors.websocket import WebSocketSensor +from airflow.providers.standard.triggers.websocket import WebSocketTrigger + +from tests_common.test_utils.version_compat import timezone + +URL = "ws://example.com/socket" +DEFAULT_DATE = timezone.datetime(2015, 1, 1) + + +class TestWebSocketSensor: + @classmethod + def setup_class(cls): + args = {"owner": "airflow", "start_date": DEFAULT_DATE} + cls.dag = DAG("test_websocket_sensor", schedule=None, default_args=args) + + @mock.patch("airflow.providers.standard.sensors.websocket.connect", autospec=True) + def test_poke_returns_true_on_message(self, mock_connect): + mock_websocket = mock.MagicMock(spec=ClientConnection) + mock_websocket.recv.return_value = "pong" + mock_connect.return_value.__enter__.return_value = mock_websocket + + sensor = WebSocketSensor(task_id="poke_true", url=URL, message_to_send="ping", dag=self.dag) + assert sensor.poke(context={}) is True + mock_websocket.send.assert_called_once_with("ping") + + @mock.patch("airflow.providers.standard.sensors.websocket.connect", autospec=True) + def test_poke_returns_false_on_timeout(self, mock_connect): + mock_websocket = mock.MagicMock(spec=ClientConnection) + mock_websocket.recv.side_effect = TimeoutError() + mock_connect.return_value.__enter__.return_value = mock_websocket + + sensor = WebSocketSensor(task_id="poke_false", url=URL, dag=self.dag) + assert sensor.poke(context={}) is False + + @mock.patch("airflow.providers.standard.sensors.websocket.connect", autospec=True) + def test_poke_waits_for_the_overall_timeout_not_poke_interval(self, mock_connect): + """recv() must be bounded by the sensor's overall timeout, not poke_interval — + otherwise a single poke could block past the sensor's declared timeout before + that timeout is ever checked.""" + mock_websocket = mock.MagicMock(spec=ClientConnection) + mock_websocket.recv.return_value = "pong" + mock_connect.return_value.__enter__.return_value = mock_websocket + + sensor = WebSocketSensor( + task_id="poke_timeout_arg", url=URL, timeout=45, poke_interval=5, dag=self.dag + ) + assert sensor.poke(context={}) is True + mock_websocket.recv.assert_called_once_with(timeout=45) + + def test_reschedule_mode_not_allowed(self): + with pytest.raises(ValueError, match="Cannot set mode to 'reschedule'. Only 'poke' is acceptable"): + WebSocketSensor(task_id="reschedule", url=URL, mode="reschedule", dag=self.dag) + + def test_task_defer_does_not_poke_first(self): + """The deferrable path must defer immediately: poke() consumes the connection, + so polling before deferring would send message_to_send and lose the reply the + trigger is supposed to wait for.""" + sensor = WebSocketSensor(task_id="defer", url=URL, deferrable=True, dag=self.dag) + + with mock.patch.object(WebSocketSensor, "poke") as mock_poke: + with pytest.raises(TaskDeferred) as exc: + sensor.execute({}) + + mock_poke.assert_not_called() + assert isinstance(exc.value.trigger, WebSocketTrigger) + assert exc.value.trigger.url == URL + + def test_execute_sync_calls_poke_exactly_once(self): + """Since poke() already blocks for the full sensor timeout, execute() must never + call it a second time — a second call would open a new connection and re-send + message_to_send.""" + sensor = WebSocketSensor(task_id="sync_timeout", url=URL, timeout=0, dag=self.dag) + + with mock.patch.object(WebSocketSensor, "poke", return_value=False) as mock_poke: Review Comment: diito -- 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. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
