Lee-W commented on code in PR #72151:
URL: https://github.com/apache/airflow/pull/72151#discussion_r4034510636
##########
providers/openai/tests/unit/openai/operators/test_openai.py:
##########
@@ -94,24 +104,103 @@ def test_openai_response_operator_execute():
response_kwargs={"instructions": "Be concise.",
"previous_response_id": "resp_prev"},
)
mock_hook_instance = Mock(spec=OpenAIHook)
- mock_hook_instance.create_response.return_value = Mock(
- spec=Response, output_text="haiku text", id="resp_123",
status="completed"
+ usage = ResponseUsage(
+ input_tokens=5,
+ input_tokens_details=InputTokensDetails(cached_tokens=1,
cache_write_tokens=0),
+ output_tokens=7,
+ output_tokens_details=OutputTokensDetails(reasoning_tokens=2),
+ total_tokens=12,
)
+ mock_response = Mock(
+ spec=Response, output_text="haiku text", id="resp_123",
status="completed", usage=usage
+ )
+ mock_hook_instance.create_response.return_value = mock_response
operator.hook = mock_hook_instance
- result = operator.execute(Context())
+ context = _build_execute_context()
+ result = operator.execute(context)
+ # Backward compat: the return value is still the aggregated output text,
unchanged
+ # by the new XCom pushes below.
assert result == "haiku text"
mock_hook_instance.create_response.assert_called_once_with(
input="Write a haiku.",
model="test_model",
instructions="Be concise.",
previous_response_id="resp_prev",
)
+ context["ti"].xcom_push.assert_any_call(key="response_id",
value="resp_123")
+ context["ti"].xcom_push.assert_any_call(
+ key="usage",
+ value={
+ "input_tokens": 5,
+ "input_tokens_details": {"cache_write_tokens": 0, "cached_tokens":
1},
+ "output_tokens": 7,
+ "output_tokens_details": {"reasoning_tokens": 2},
+ "total_tokens": 12,
+ },
+ )
+
+
+def test_openai_response_operator_execute_without_usage():
+ operator = OpenAIResponseOperator(
+ task_id=TASK_ID, conn_id=CONN_ID, input_text="Write a haiku.",
model="test_model"
+ )
+ mock_hook_instance = Mock(spec=OpenAIHook)
+ mock_response = Mock(
+ spec=Response, output_text="haiku text", id="resp_123",
status="completed", usage=None
+ )
+ mock_hook_instance.create_response.return_value = mock_response
+ operator.hook = mock_hook_instance
+
+ context = _build_execute_context()
+ result = operator.execute(context)
+
+ assert result == "haiku text"
+ context["ti"].xcom_push.assert_any_call(key="usage", value=None)
+
+
+def test_openai_response_operator_execute_skips_xcom_push_when_disabled():
Review Comment:
Merged enabled/disabled `do_xcom_push` tests into one parametrized test
asserting `xcom_push.call_count`, and the disabled case's mock response now
sets a real `usage` value via the shared `_build_completed_response()` helper.
--
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]
