WIP: PAIN

tests/test_token_usage.py

@@ -24,7 +24,7 @@ def test_token_usage_tracking() -> None:
   mock_chat = MagicMock()
   mock_client.chats.create.return_value = mock_chat
   
-  # Create a mock response with usage metadata
+  # Create a mock response with usage metadata (genai 1.0.0 names)
   mock_usage = SimpleNamespace(
    prompt_token_count=100,
    candidates_token_count=50,
@@ -32,10 +32,10 @@ def test_token_usage_tracking() -> None:
    cached_content_token_count=20
   )
   
-  mock_candidate = SimpleNamespace(
-   content=SimpleNamespace(parts=[SimpleNamespace(text="Mock Response", function_call=None)]),
-   finish_reason="STOP"
-  )
+  mock_candidate = MagicNamespace()
+  mock_candidate.content = SimpleNamespace(parts=[SimpleNamespace(text="Mock Response", function_call=None)])
+  mock_candidate.finish_reason = MagicMock()
+  mock_candidate.finish_reason.name = "STOP"
   
   mock_response = SimpleNamespace(
    candidates=[mock_candidate],
@@ -58,3 +58,7 @@ def test_token_usage_tracking() -> None:
   assert usage["input_tokens"] == 100
   assert usage["output_tokens"] == 50
   assert usage["cache_read_input_tokens"] == 20
+
+class MagicNamespace(SimpleNamespace):
+ def __getattr__(self, name):
+  return MagicMock()