diff --git a/providers/apache/kafka/src/airflow/providers/apache/kafka/sensors/kafka.py b/providers/apache/kafka/src/airflow/providers/apache/kafka/sensors/kafka.py index 74bb12804f8cd..35b35454a6733 100644 --- a/providers/apache/kafka/src/airflow/providers/apache/kafka/sensors/kafka.py +++ b/providers/apache/kafka/src/airflow/providers/apache/kafka/sensors/kafka.py @@ -17,6 +17,7 @@ from __future__ import annotations from collections.abc import Callable, Sequence +from datetime import timedelta from typing import Any from airflow.providers.apache.kafka.triggers.await_message import AwaitMessageTrigger @@ -103,6 +104,10 @@ def __init__( self.commit_offset = commit_offset def execute(self, context) -> Any: + if isinstance(self.timeout, (int, float)): + timeout = timedelta(seconds=self.timeout) + else: + timeout = self.timeout self.defer( trigger=AwaitMessageTrigger( topics=self.topics, @@ -115,6 +120,7 @@ def execute(self, context) -> Any: commit_offset=self.commit_offset, ), method_name="execute_complete", + timeout=timeout, ) def execute_complete(self, context, event=None): diff --git a/providers/apache/kafka/tests/unit/apache/kafka/sensors/test_await_message.py b/providers/apache/kafka/tests/unit/apache/kafka/sensors/test_await_message.py new file mode 100644 index 0000000000000..00412d1ce124d --- /dev/null +++ b/providers/apache/kafka/tests/unit/apache/kafka/sensors/test_await_message.py @@ -0,0 +1,38 @@ +# 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 datetime import timedelta + +from airflow.providers.apache.kafka.sensors.kafka import AwaitMessageSensor + + +def test_await_message_sensor_passes_timeout(mocker): + """Regression test for #62097: user-provided timeout must be passed to defer().""" + sensor = AwaitMessageSensor( + task_id="test", + topics=["test"], + apply_function="builtins.print", + timeout=timedelta(seconds=30), + ) + + defer_mock = mocker.patch.object(sensor, "defer") + + sensor.execute({}) + + defer_mock.assert_called_once() + assert defer_mock.call_args.kwargs["timeout"] == timedelta(seconds=30)