Lee-W commented on code in PR #72151:
URL: https://github.com/apache/airflow/pull/72151#discussion_r4034456220
##########
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():
+ operator = OpenAIResponseOperator(
+ task_id=TASK_ID,
+ conn_id=CONN_ID,
+ input_text="Write a haiku.",
+ model="test_model",
+ do_xcom_push=False,
+ )
+ mock_hook_instance = Mock(spec=OpenAIHook)
+ mock_response = Mock(spec=Response, output_text="haiku text",
id="resp_123", status="completed")
+ 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_not_called()
+
+
+def test_openai_response_operator_templates_input_text_and_response_kwargs():
Review Comment:
Added parametrize cases for `{% raw %}` escape and nested tool-schema
rendering in
`test_openai_response_operator_templates_input_text_and_response_kwargs`. The
nested case also asserts the non-template `max_retries` int survives unchanged.
--
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]
