diff --git a/providers/openai/docs/changelog.rst b/providers/openai/docs/changelog.rst index cc2a0bd901fe3..742450896dd34 100644 --- a/providers/openai/docs/changelog.rst +++ b/providers/openai/docs/changelog.rst @@ -20,6 +20,25 @@ Changelog --------- +.. note:: + A deferred ``OpenAITriggerBatchOperator`` that runs out of patience now raises + ``OpenAIBatchTimeout``, matching the exception the non-deferrable path has always raised + for the same condition. Previously every non-success outcome of a deferred batch raised + the same ``OpenAIBatchJobException``, so a timeout and a genuine batch failure could not + be told apart or handled separately. + + A cancelled batch now raises ``OpenAIBatchCancelled``, a subclass of + ``OpenAIBatchJobException``, so existing code that catches ``OpenAIBatchJobException`` + keeps working unchanged while callers that want to distinguish cancellation can catch the + subclass specifically. + +.. note:: + A deferred ``OpenAITriggerBatchOperator`` that times out now requests cancellation of the + batch, matching the non-deferrable path. Previously a deferred timeout only failed the + task and left the batch running (and billing) on OpenAI's side. Cancellation on OpenAI's + side is asynchronous, so the batch reports ``cancelling`` for a while before it settles as + ``cancelled``. + 1.8.2 ..... diff --git a/providers/openai/src/airflow/providers/openai/exceptions.py b/providers/openai/src/airflow/providers/openai/exceptions.py index 85f015880c548..4b2b2c7db0d21 100644 --- a/providers/openai/src/airflow/providers/openai/exceptions.py +++ b/providers/openai/src/airflow/providers/openai/exceptions.py @@ -24,6 +24,17 @@ class OpenAIBatchJobException(AirflowException): """Raise when OpenAI Batch Job fails to start AFTER processing the request.""" +class OpenAIBatchCancelled(OpenAIBatchJobException): + """ + Raise when an OpenAI Batch Job was cancelled. + + Cancellation is a decision, not a failure, so it gets its own subclass: callers + that want to distinguish "someone cancelled this batch" from "the batch failed" + can catch this specifically, while existing handlers written against + ``OpenAIBatchJobException`` keep working unchanged. + """ + + class OpenAIBatchTimeout(AirflowException): """Raise when OpenAI Batch Job times out.""" diff --git a/providers/openai/src/airflow/providers/openai/hooks/openai.py b/providers/openai/src/airflow/providers/openai/hooks/openai.py index 97dcc4ceec8dc..0a8dba11af03e 100644 --- a/providers/openai/src/airflow/providers/openai/hooks/openai.py +++ b/providers/openai/src/airflow/providers/openai/hooks/openai.py @@ -54,8 +54,9 @@ from openai.types.vector_stores import VectorStoreFile, VectorStoreFileBatch, VectorStoreFileDeleted from airflow.exceptions import AirflowProviderDeprecationWarning from airflow.providers.common.compat.module_loading import import_string -from airflow.providers.common.compat.sdk import BaseHook +from airflow.providers.common.compat.sdk import AirflowException, BaseHook from airflow.providers.openai.exceptions import ( + OpenAIBatchCancelled, OpenAIBatchJobException, OpenAIBatchTimeout, OpenAITriggerEventError, @@ -95,6 +96,29 @@ def is_in_progress(cls, status: str) -> bool: #: Statuses the provider's trigger emits in its terminal event. TRIGGER_EVENT_STATUSES = frozenset({"success", "error", "cancelled"}) +#: Maps the trigger's ``termination_reason`` field to the exception ``execute_complete`` +#: should raise. Keyed on the reason field, never on the message text, so that a +#: rewording of the trigger's message never silently changes which exception a +#: downstream task can catch. +_TERMINATION_REASON_EXCEPTIONS: dict[str, type[AirflowException]] = { + "timeout": OpenAIBatchTimeout, + "cancelled": OpenAIBatchCancelled, +} + + +def build_batch_error(message: str, termination_reason: str | None) -> AirflowException: + """ + Build (but do not raise) the exception matching a trigger event's termination reason. + + ``termination_reason`` is ``None`` when the event was produced by a trigger + serialized before this field existed (a rolling upgrade in flight); that case + falls back to ``OpenAIBatchJobException``, matching today's behavior. + """ + if termination_reason is None: + return OpenAIBatchJobException(message) + exception_class = _TERMINATION_REASON_EXCEPTIONS.get(termination_reason, OpenAIBatchJobException) + return exception_class(message) + def validate_execute_complete_event(event: dict[str, Any] | None = None) -> dict[str, Any]: """ @@ -670,10 +694,10 @@ def wait_for_batch(self, batch_id: str, wait_seconds: float = 3, timeout: float if batch.status == BatchStatus.FAILED: raise OpenAIBatchJobException(f"Batch failed - \n{batch_id}") if batch.status in (BatchStatus.CANCELLED, BatchStatus.CANCELLING): - raise OpenAIBatchJobException(f"Batch failed - batch was cancelled:\n{batch_id}") + raise OpenAIBatchCancelled(f"Batch failed - batch was cancelled:\n{batch_id}") if batch.status == BatchStatus.EXPIRED: raise OpenAIBatchJobException( - f"Batch failed - batch couldn't be completed within the hour time window :\n{batch_id}" + f"Batch failed - batch couldn't be completed within its completion window:\n{batch_id}" ) raise OpenAIBatchJobException( diff --git a/providers/openai/src/airflow/providers/openai/operators/openai.py b/providers/openai/src/airflow/providers/openai/operators/openai.py index dfed6e48d51a9..57b0c0c365f15 100644 --- a/providers/openai/src/airflow/providers/openai/operators/openai.py +++ b/providers/openai/src/airflow/providers/openai/operators/openai.py @@ -22,8 +22,11 @@ from typing import TYPE_CHECKING, Any, Literal from airflow.providers.common.compat.sdk import BaseOperator, conf -from airflow.providers.openai.exceptions import OpenAIBatchJobException -from airflow.providers.openai.hooks.openai import OpenAIHook, validate_execute_complete_event +from airflow.providers.openai.hooks.openai import ( + OpenAIHook, + build_batch_error, + validate_execute_complete_event, +) from airflow.providers.openai.triggers.openai import OpenAIBatchTrigger if TYPE_CHECKING: @@ -145,10 +148,26 @@ class OpenAITriggerBatchOperator(BaseOperator): :param wait_seconds: Optional. Number of seconds between checks. Only used when ``deferrable`` is False. Defaults to 3 seconds. :param timeout: Optional. The amount of time, in seconds, to wait for the request to complete. - Only used when ``deferrable`` is False. Defaults to 24 hour, which is the SLA for OpenAI Batch API. + Used in both modes: in the synchronous path it bounds ``wait_for_batch``; in the + deferrable path it bounds the trigger's poll loop. When the deferrable path times out, + the operator requests cancellation of the batch using the batch id carried by the + trigger event, mirroring the synchronous path. Cancellation on OpenAI's side is + asynchronous — the batch reports ``cancelling`` for up to 10 minutes before it settles + as ``cancelled`` — so this only *requests* cancellation, it does not wait for it. If + ``execution_timeout`` is set shorter than ``timeout``, the scheduler's deferral timeout + fires first: the task is failed with ``TaskDeferralTimeout`` before the trigger ever + times out, ``execute_complete`` is never called, and this cancellation path does not + run. Defaults to 24 hour, which is the SLA for OpenAI Batch API. :param wait_for_completion: Optional. Whether to wait for the batch to complete. If set to False, the operator will return immediately after triggering the batch. Defaults to True. + When ``deferrable`` is True and the batch does not reach a terminal state, ``execute_complete`` + raises :class:`~airflow.providers.openai.exceptions.OpenAIBatchTimeout`, matching the exception + raised by the synchronous path for the same condition. A cancelled batch raises + :class:`~airflow.providers.openai.exceptions.OpenAIBatchCancelled` (a subclass of + :class:`~airflow.providers.openai.exceptions.OpenAIBatchJobException`), and any other failure + raises :class:`~airflow.providers.openai.exceptions.OpenAIBatchJobException`. + .. seealso:: For more information on how to use this operator, please take a look at the guide: :ref:`howto/operator:OpenAITriggerBatchOperator` @@ -207,15 +226,58 @@ def execute_complete(self, context: Context, event: Any = None) -> str: Invoke this callback when the trigger fires; return immediately. Relies on trigger to throw an exception, otherwise it assumes execution was - successful. + successful. The exception raised depends on the event's ``termination_reason``: + ``OpenAIBatchTimeout`` for a timeout, ``OpenAIBatchCancelled`` for a cancellation, + and ``OpenAIBatchJobException`` for any other failure (including events from a + trigger serialized before ``termination_reason`` existed). + + On a timeout, cancellation of the batch is requested before the timeout is raised + (see :meth:`_cancel_batch_quietly`). No other termination reason triggers + cancellation: a ``polling_error`` may be a transient, Airflow-side failure rather than + a real batch problem, and cancellation is irreversible, so it is left alone to run to + its own 24-hour completion window instead. """ event = validate_execute_complete_event(event) if event["status"] != "success": - raise OpenAIBatchJobException(event["message"]) + if event.get("termination_reason") == "timeout": + batch_id = event.get("batch_id") + if batch_id: + self.log.warning( + "%s timed out waiting for batch %s; requesting cancellation.", + self.task_id, + batch_id, + ) + self._cancel_batch_quietly(batch_id) + else: + self.log.warning( + "%s timed out but the trigger event carried no batch_id; " + "skipping cancellation request.", + self.task_id, + ) + raise build_batch_error(event["message"], event.get("termination_reason")) self.log.info("%s completed successfully.", self.task_id) return event["batch_id"] + def _cancel_batch_quietly(self, batch_id: str) -> None: + """ + Best-effort request to cancel a batch; never raises. + + Called from ``execute_complete`` after a deferred timeout, using the batch id carried + by the trigger event rather than ``self.batch_id`` — this method runs on a resumed task + instance, a fresh operator object on which ``execute``'s assignment to ``self.batch_id`` + never happened, so ``self.batch_id`` is ``None`` here. + + Cancellation on OpenAI's side is asynchronous: the batch reports ``cancelling`` for up + to 10 minutes before it settles as ``cancelled``, so this only requests cancellation. A + failure to cancel is logged, not raised, so it never masks the timeout that is the + task's real failure reason. + """ + try: + self.hook.cancel_batch(batch_id) + except Exception as e: + self.log.warning("Failed to request cancellation of batch %s: %s", batch_id, e) + def on_kill(self) -> None: """Cancel the batch if task is cancelled.""" if self.batch_id: diff --git a/providers/openai/src/airflow/providers/openai/triggers/openai.py b/providers/openai/src/airflow/providers/openai/triggers/openai.py index 49fd0900cc022..767aaa0c4bc41 100644 --- a/providers/openai/src/airflow/providers/openai/triggers/openai.py +++ b/providers/openai/src/airflow/providers/openai/triggers/openai.py @@ -103,6 +103,7 @@ async def run(self) -> AsyncIterator[TriggerEvent]: yield TriggerEvent( { "status": "error", + "termination_reason": "timeout", "message": ( f"Batch {self.batch_id} has not reached a terminal status after " f"{elapsed:.0f} seconds." @@ -116,6 +117,7 @@ async def run(self) -> AsyncIterator[TriggerEvent]: yield TriggerEvent( { "status": "success", + "termination_reason": "completed", "message": f"Batch {self.batch_id} has completed successfully.", "batch_id": self.batch_id, } @@ -124,6 +126,7 @@ async def run(self) -> AsyncIterator[TriggerEvent]: yield TriggerEvent( { "status": "cancelled", + "termination_reason": "cancelled", "message": f"Batch {self.batch_id} has been cancelled.", "batch_id": self.batch_id, } @@ -132,6 +135,7 @@ async def run(self) -> AsyncIterator[TriggerEvent]: yield TriggerEvent( { "status": "error", + "termination_reason": "failed", "message": f"Batch failed:\n{self.batch_id}", "batch_id": self.batch_id, } @@ -140,7 +144,8 @@ async def run(self) -> AsyncIterator[TriggerEvent]: yield TriggerEvent( { "status": "error", - "message": f"Batch couldn't be completed within the hour time window :\n{self.batch_id}", + "termination_reason": "expired", + "message": f"Batch couldn't be completed within its completion window:\n{self.batch_id}", "batch_id": self.batch_id, } ) @@ -148,9 +153,17 @@ async def run(self) -> AsyncIterator[TriggerEvent]: yield TriggerEvent( { "status": "error", + "termination_reason": "unexpected_status", "message": f"Batch {self.batch_id} has failed.", "batch_id": self.batch_id, } ) except Exception as e: - yield TriggerEvent({"status": "error", "message": str(e), "batch_id": self.batch_id}) + yield TriggerEvent( + { + "status": "error", + "termination_reason": "polling_error", + "message": str(e), + "batch_id": self.batch_id, + } + ) diff --git a/providers/openai/tests/unit/openai/hooks/test_openai.py b/providers/openai/tests/unit/openai/hooks/test_openai.py index d132a6a3d9582..155e45c46c68e 100644 --- a/providers/openai/tests/unit/openai/hooks/test_openai.py +++ b/providers/openai/tests/unit/openai/hooks/test_openai.py @@ -39,6 +39,7 @@ from airflow.exceptions import AirflowProviderDeprecationWarning from airflow.models import Connection from airflow.providers.openai.exceptions import ( + OpenAIBatchCancelled, OpenAIBatchJobException, OpenAIBatchTimeout, OpenAITriggerEventError, @@ -644,6 +645,24 @@ def test_wait_for_in_progress_batch_timeout(mock_openai_hook, mock_wip_batch): assert mock_openai_hook.conn.batches.cancel.call_count == 1 +@pytest.mark.parametrize("status", ["cancelled", "cancelling"]) +def test_wait_for_cancelled_batch_raises_exact_cancelled_type(mock_openai_hook, status): + """``OpenAIBatchCancelled`` is a subclass of ``OpenAIBatchJobException``, so asserting + only the base class would stay green even if this raised the wrong (base) type. Assert + the exact type to prove the exception was actually narrowed. + """ + mock_openai_hook.conn.batches.retrieve.return_value = create_batch(status) + with pytest.raises(OpenAIBatchCancelled): + mock_openai_hook.wait_for_batch(batch_id=BATCH_ID) + + +def test_wait_for_expired_batch_message_does_not_mention_hour_window(mock_openai_hook): + mock_openai_hook.conn.batches.retrieve.return_value = create_batch("expired") + with pytest.raises(OpenAIBatchJobException, match="completion window") as exc_info: + mock_openai_hook.wait_for_batch(batch_id=BATCH_ID) + assert "hour time window" not in str(exc_info.value) + + def test_openai_hook_test_connection(mock_openai_hook): result, message = mock_openai_hook.test_connection() assert result is True diff --git a/providers/openai/tests/unit/openai/operators/test_openai.py b/providers/openai/tests/unit/openai/operators/test_openai.py index 0ec78bb182c51..556cecb3e7698 100644 --- a/providers/openai/tests/unit/openai/operators/test_openai.py +++ b/providers/openai/tests/unit/openai/operators/test_openai.py @@ -23,7 +23,12 @@ from openai.types.responses import Response from airflow.providers.common.compat.sdk import Context, TaskDeferred -from airflow.providers.openai.exceptions import OpenAIBatchJobException, OpenAITriggerEventError +from airflow.providers.openai.exceptions import ( + OpenAIBatchCancelled, + OpenAIBatchJobException, + OpenAIBatchTimeout, + OpenAITriggerEventError, +) from airflow.providers.openai.hooks.openai import OpenAIHook from airflow.providers.openai.operators.openai import ( OpenAIEmbeddingOperator, @@ -187,3 +192,115 @@ def test_failed_event_raises(self, event): def test_invalid_event_raises_instead_of_succeeding(self, event): with pytest.raises(OpenAITriggerEventError): self._operator().execute_complete(Context(), event) + + @pytest.mark.parametrize( + ("termination_reason", "expected_exc"), + [ + pytest.param("timeout", OpenAIBatchTimeout, id="timeout"), + pytest.param("cancelled", OpenAIBatchCancelled, id="cancelled"), + pytest.param("failed", OpenAIBatchJobException, id="failed"), + pytest.param("expired", OpenAIBatchJobException, id="expired"), + pytest.param("unexpected_status", OpenAIBatchJobException, id="unexpected-status"), + pytest.param("polling_error", OpenAIBatchJobException, id="polling-error"), + ], + ) + def test_execute_complete_raises_exception_matching_termination_reason( + self, termination_reason, expected_exc + ): + event = { + "status": "error", + "termination_reason": termination_reason, + "message": "boom", + "batch_id": BATCH_ID, + } + with pytest.raises(expected_exc, match="boom"): + self._operator().execute_complete(Context(), event) + + @pytest.mark.parametrize("status", ["error", "cancelled"]) + def test_execute_complete_missing_termination_reason_falls_back(self, status): + """A trigger serialized before ``termination_reason`` existed sends an event without + that key; ``execute_complete`` must fall back to ``OpenAIBatchJobException`` exactly, + not raise ``KeyError``. + """ + event = {"status": status, "message": "boom", "batch_id": BATCH_ID} + with pytest.raises(OpenAIBatchJobException, match="boom") as exc_info: + self._operator().execute_complete(Context(), event) + assert type(exc_info.value) is OpenAIBatchJobException + + def test_timeout_requests_cancellation_using_event_batch_id(self): + """This is a regression lock: the resumed task is a fresh operator instance, so + ``self.batch_id`` is ``None`` here (``execute``'s assignment never happened on this + object). Cancellation must use ``event["batch_id"]``; if this test is made to pass by + reading ``self.batch_id`` instead, it should fail again as soon as that read returns + ``None`` for a real resumed task. + """ + operator = self._operator() + assert operator.batch_id is None + mock_hook_instance = Mock(spec=OpenAIHook) + operator.hook = mock_hook_instance + event = { + "status": "error", + "termination_reason": "timeout", + "message": "boom", + "batch_id": BATCH_ID, + } + + with pytest.raises(OpenAIBatchTimeout): + operator.execute_complete(Context(), event) + + mock_hook_instance.cancel_batch.assert_called_once_with(BATCH_ID) + + def test_cancel_failure_does_not_mask_timeout(self): + operator = self._operator() + mock_hook_instance = Mock(spec=OpenAIHook) + mock_hook_instance.cancel_batch.side_effect = RuntimeError("cancel failed") + operator.hook = mock_hook_instance + event = { + "status": "error", + "termination_reason": "timeout", + "message": "boom", + "batch_id": BATCH_ID, + } + + with pytest.raises(OpenAIBatchTimeout): + operator.execute_complete(Context(), event) + + mock_hook_instance.cancel_batch.assert_called_once_with(BATCH_ID) + + @pytest.mark.parametrize( + "termination_reason", + [ + "failed", + "cancelled", + "expired", + "polling_error", + "unexpected_status", + None, # a trigger serialized before `termination_reason` existed sends no such key + ], + ) + def test_non_timeout_termination_reasons_do_not_cancel(self, termination_reason): + operator = self._operator() + mock_hook_instance = Mock(spec=OpenAIHook) + operator.hook = mock_hook_instance + event = {"status": "error", "message": "boom", "batch_id": BATCH_ID} + if termination_reason is not None: + event["termination_reason"] = termination_reason + + with pytest.raises(OpenAIBatchJobException): + operator.execute_complete(Context(), event) + + mock_hook_instance.cancel_batch.assert_not_called() + + def test_timeout_with_no_batch_id_skips_cancellation(self): + """A timeout event that (hypothetically) carries no ``batch_id`` must not crash and + must not call ``cancel_batch`` — but the original timeout error still has to surface. + """ + operator = self._operator() + mock_hook_instance = Mock(spec=OpenAIHook) + operator.hook = mock_hook_instance + event = {"status": "error", "termination_reason": "timeout", "message": "boom"} + + with pytest.raises(OpenAIBatchTimeout): + operator.execute_complete(Context(), event) + + mock_hook_instance.cancel_batch.assert_not_called() diff --git a/providers/openai/tests/unit/openai/test_exceptions.py b/providers/openai/tests/unit/openai/test_exceptions.py index fabaad35343f0..b9c6a14e44667 100644 --- a/providers/openai/tests/unit/openai/test_exceptions.py +++ b/providers/openai/tests/unit/openai/test_exceptions.py @@ -21,7 +21,11 @@ import pytest -from airflow.providers.openai.exceptions import OpenAIBatchJobException, OpenAIBatchTimeout +from airflow.providers.openai.exceptions import ( + OpenAIBatchCancelled, + OpenAIBatchJobException, + OpenAIBatchTimeout, +) from airflow.providers.openai.hooks.openai import OpenAIHook @@ -30,6 +34,7 @@ [ OpenAIBatchTimeout, OpenAIBatchJobException, + OpenAIBatchCancelled, ], ) def test_wait_for_batch_raise_exception(exception_class): @@ -38,3 +43,11 @@ def test_wait_for_batch_raise_exception(exception_class): hook = mock_hook_instance with pytest.raises(exception_class): hook.wait_for_batch(batch_id="batch_id") + + +def test_batch_cancelled_is_subclass_of_batch_job_exception(): + """Cancellation is deliberately a subclass, not a sibling, of the generic batch failure + exception: existing ``except OpenAIBatchJobException`` handlers must keep working + unchanged after cancellation gets its own exception type. + """ + assert issubclass(OpenAIBatchCancelled, OpenAIBatchJobException) diff --git a/providers/openai/tests/unit/openai/triggers/test_openai.py b/providers/openai/tests/unit/openai/triggers/test_openai.py index 0a2c4322a5ed6..d72620ed3c501 100644 --- a/providers/openai/tests/unit/openai/triggers/test_openai.py +++ b/providers/openai/tests/unit/openai/triggers/test_openai.py @@ -119,22 +119,28 @@ def test_rejects_both_timeout_and_end_time(self): @pytest.mark.asyncio @pytest.mark.parametrize( - ("mock_batch_status", "mock_status", "mock_message"), + ("mock_batch_status", "mock_status", "mock_termination_reason", "mock_message"), [ - (str(BatchStatus.COMPLETED), "success", "Batch batch_id has completed successfully."), - (str(BatchStatus.CANCELLING), "cancelled", "Batch batch_id has been cancelled."), - (str(BatchStatus.CANCELLED), "cancelled", "Batch batch_id has been cancelled."), - (str(BatchStatus.FAILED), "error", "Batch failed:\nbatch_id"), + ( + str(BatchStatus.COMPLETED), + "success", + "completed", + "Batch batch_id has completed successfully.", + ), + (str(BatchStatus.CANCELLING), "cancelled", "cancelled", "Batch batch_id has been cancelled."), + (str(BatchStatus.CANCELLED), "cancelled", "cancelled", "Batch batch_id has been cancelled."), + (str(BatchStatus.FAILED), "error", "failed", "Batch failed:\nbatch_id"), ( str(BatchStatus.EXPIRED), "error", - "Batch couldn't be completed within the hour time window :\nbatch_id", + "expired", + "Batch couldn't be completed within its completion window:\nbatch_id", ), ], ) @mock.patch("airflow.providers.openai.hooks.openai.OpenAIHook.get_batch") async def test_openai_batch_for_terminal_status( - self, mock_batch, mock_batch_status, mock_status, mock_message + self, mock_batch, mock_batch_status, mock_status, mock_termination_reason, mock_message ): """Assert that run trigger messages in case of job finished""" mock_batch.return_value = self.mock_get_batch(mock_batch_status) @@ -146,6 +152,7 @@ async def test_openai_batch_for_terminal_status( ) expected_result = { "status": mock_status, + "termination_reason": mock_termination_reason, "message": mock_message, "batch_id": self.BATCH_ID, } @@ -187,6 +194,7 @@ async def test_openai_batch_for_timeout(self, mock_monotonic, mock_batch, mock_b await asyncio.sleep(0.1) event = task.result() assert event.payload["status"] == "error" + assert event.payload["termination_reason"] == "timeout" assert f"Batch {self.BATCH_ID} has not reached a terminal status after" in event.payload["message"] asyncio.get_event_loop().stop() @@ -235,6 +243,7 @@ async def test_openai_batch_yields_single_terminal_event(self, mock_batch): TriggerEvent( { "status": "success", + "termination_reason": "completed", "message": f"Batch {self.BATCH_ID} has completed successfully.", "batch_id": self.BATCH_ID, } @@ -254,6 +263,7 @@ async def test_openai_batch_for_unexpected_error(self, mock_batch): ) expected_result = { "status": "error", + "termination_reason": "polling_error", "message": "'float' object has no attribute 'status'", "batch_id": self.BATCH_ID, } @@ -261,3 +271,26 @@ async def test_openai_batch_for_unexpected_error(self, mock_batch): await asyncio.sleep(0.1) assert TriggerEvent(expected_result) == task.result() asyncio.get_event_loop().stop() + + @pytest.mark.asyncio + @mock.patch("airflow.providers.openai.hooks.openai.OpenAIHook.get_batch") + async def test_openai_batch_for_unexpected_status(self, mock_batch): + """A batch status outside the known terminal set falls into the `unexpected_status` branch.""" + mock_batch.return_value = self.mock_get_batch("validating") + mock_batch.return_value.status = "some_future_status" + trigger = OpenAIBatchTrigger( + conn_id=self.CONN_ID, + batch_id=self.BATCH_ID, + poll_interval=self.POLL_INTERVAL, + timeout=self.TIMEOUT, + ) + expected_result = { + "status": "error", + "termination_reason": "unexpected_status", + "message": f"Batch {self.BATCH_ID} has failed.", + "batch_id": self.BATCH_ID, + } + task = asyncio.create_task(trigger.run().__anext__()) + await asyncio.sleep(0.1) + assert TriggerEvent(expected_result) == task.result() + asyncio.get_event_loop().stop()