kaxil commented on code in PR #71463:
URL: https://github.com/apache/airflow/pull/71463#discussion_r3766978085
##########
providers/anthropic/tests/unit/anthropic/hooks/test_anthropic.py:
##########
@@ -308,6 +315,94 @@ def test_drops_a_none_budget(self):
assert "budget" not in client.beta.sessions.create.call_args.kwargs
+class TestGetSessionUsage:
+ """
+ Built on the real SDK usage models rather than mocks.
+
+ Constructing the real model is not on its own a check on field names: it
sets
+ ``extra="allow"``, so a misspelling is accepted as an extra field rather
than rejected.
+ ``test_usage_fields_exist_on_the_sdk_model`` is what pins the names to the
declared
+ schema. Only the session wrapper is a stand-in, since
``get_session_usage`` reads
+ nothing from it but ``.usage``.
+ """
+
+ def test_usage_fields_exist_on_the_sdk_model(self):
+ # The reported breakdown must be fields the SDK actually declares.
active_seconds
+ # lived on session.stats rather than session.usage as recently as
0.117.0, and the
+ # floor is an open-ended >= on a beta API -- if a field moves again,
summarize_usage
+ # would return it as None (or raise, swallowed into a warning) with no
other signal.
+ assert set(BetaManagedAgentsSessionUsage.model_fields) >= {
+ "input_tokens",
+ "output_tokens",
+ "cache_read_input_tokens",
+ "cache_creation",
+ "server_tool_use",
+ "active_seconds",
+ "list_cost",
+ }
+
+ @staticmethod
+ def _session(list_cost, cache_creation=None, server_tool_use=None):
Review Comment:
Applied. It also disambiguates it from the module-level `_session(status,
...)` helper further up the file, which is a different thing entirely.
##########
providers/anthropic/tests/unit/anthropic/operators/test_agent.py:
##########
@@ -34,8 +34,10 @@
pytest.importorskip("anthropic")
-def _context():
- return {"ti": mock.MagicMock()}
+def _context(try_number=1):
Review Comment:
Applied, across all 30 call sites in the file.
##########
providers/anthropic/tests/unit/anthropic/hooks/test_anthropic.py:
##########
@@ -308,6 +315,94 @@ def test_drops_a_none_budget(self):
assert "budget" not in client.beta.sessions.create.call_args.kwargs
+class TestGetSessionUsage:
+ """
+ Built on the real SDK usage models rather than mocks.
+
+ Constructing the real model is not on its own a check on field names: it
sets
+ ``extra="allow"``, so a misspelling is accepted as an extra field rather
than rejected.
+ ``test_usage_fields_exist_on_the_sdk_model`` is what pins the names to the
declared
+ schema. Only the session wrapper is a stand-in, since
``get_session_usage`` reads
+ nothing from it but ``.usage``.
+ """
+
+ def test_usage_fields_exist_on_the_sdk_model(self):
+ # The reported breakdown must be fields the SDK actually declares.
active_seconds
+ # lived on session.stats rather than session.usage as recently as
0.117.0, and the
+ # floor is an open-ended >= on a beta API -- if a field moves again,
summarize_usage
+ # would return it as None (or raise, swallowed into a warning) with no
other signal.
+ assert set(BetaManagedAgentsSessionUsage.model_fields) >= {
+ "input_tokens",
+ "output_tokens",
+ "cache_read_input_tokens",
+ "cache_creation",
+ "server_tool_use",
+ "active_seconds",
+ "list_cost",
+ }
+
+ @staticmethod
+ def _session(list_cost, cache_creation=None, server_tool_use=None):
+ usage = BetaManagedAgentsSessionUsage(
+ input_tokens=827,
+ output_tokens=17065,
+ cache_read_input_tokens=0,
+ active_seconds=91.2,
+ list_cost=list_cost,
+ cache_creation=cache_creation,
+ server_tool_use=server_tool_use,
+ )
+ return SimpleNamespace(usage=usage)
+
+ def test_flattens_to_json_safe_scalars(self):
+ hook, client = _make_hook()
+ client.beta.sessions.retrieve.return_value = self._session(
+ BetaMonetaryAmount(amount="44", currency="USD")
+ )
+ assert hook.get_session_usage("s") == {
+ "input_tokens": 827,
+ "output_tokens": 17065,
+ "cache_read_input_tokens": 0,
+ "cache_creation": None,
+ "server_tool_use": None,
+ "active_seconds": 91.2,
+ "list_cost": {"amount": "44", "currency": "USD"},
+ }
+
+ def test_reports_every_billable_dimension(self):
+ # list_cost is None exactly when a caller must price the run from
usage, so the
+ # expensive side (cache writes, server tool calls) cannot be missing.
+ hook, client = _make_hook()
+ client.beta.sessions.retrieve.return_value = self._session(
+ None,
+ cache_creation=BetaManagedAgentsCacheCreationUsage(
+ ephemeral_5m_input_tokens=120, ephemeral_1h_input_tokens=340
+ ),
+
server_tool_use=BetaManagedAgentsServerToolUsage(web_search_requests=3,
web_fetch_requests=5),
+ )
+ usage = hook.get_session_usage("s")
Review Comment:
Done -- it now asserts the whole dict. Better than my version: the per-key
asserts would not have caught an unexpected extra key, and this pins the exact
shape that lands in XCom.
--
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]
