fix: adjust some tests for earlier versions of the google genai package

Author: constantinius

tests/integrations/google_genai/test_google_genai.py

@@ -1043,7 +1043,9 @@ def test_embed_content(
         assert SPANDATA.GEN_AI_EMBEDDINGS_INPUT not in embed_span["data"]
 
     # Check usage data (sum of token counts from statistics: 10 + 15 = 25)
-    assert embed_span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 25
+    # Note: Only available in newer versions with ContentEmbeddingStatistics
+    if SPANDATA.GEN_AI_USAGE_INPUT_TOKENS in embed_span["data"]:
+        assert embed_span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 25
 
 
 def test_embed_content_string_input(sentry_init, capture_events, mock_genai_client):
@@ -1088,7 +1090,9 @@ def test_embed_content_string_input(sentry_init, capture_events, mock_genai_clie
     input_texts = json.loads(embed_span["data"][SPANDATA.GEN_AI_EMBEDDINGS_INPUT])
     assert input_texts == ["Single text input"]
     # Should use token_count from statistics (5), not billable_character_count (10)
-    assert embed_span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 5
+    # Note: Only available in newer versions with ContentEmbeddingStatistics
+    if SPANDATA.GEN_AI_USAGE_INPUT_TOKENS in embed_span["data"]:
+        assert embed_span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 5
 
 
 def test_embed_content_error_handling(sentry_init, capture_events, mock_genai_client):
@@ -1122,37 +1126,44 @@ def test_embed_content_error_handling(sentry_init, capture_events, mock_genai_cl
     assert error_event["exception"]["values"][0]["mechanism"]["type"] == "google_genai"
 
 
-def test_embed_content_without_metadata(sentry_init, capture_events, mock_genai_client):
-    """Test embed_content response without metadata (MLDev API)."""
+def test_embed_content_without_statistics(
+    sentry_init, capture_events, mock_genai_client
+):
+    """Test embed_content response without statistics (older package versions)."""
     sentry_init(
         integrations=[GoogleGenAIIntegration()],
         traces_sample_rate=1.0,
     )
     events = capture_events()
 
-    # Response without metadata (typical for MLDev, not Vertex AI)
-    mldev_response = {
+    # Response without statistics (typical for older google-genai versions)
+    # Embeddings exist but don't have the statistics field
+    old_version_response = {
         "embeddings": [
             {
                 "values": [0.1, 0.2, 0.3],
             },
+            {
+                "values": [0.2, 0.3, 0.4],
+            },
         ],
     }
-    mock_http_response = create_mock_http_response(mldev_response)
+    mock_http_response = create_mock_http_response(old_version_response)
 
     with mock.patch.object(
         mock_genai_client._api_client, "request", return_value=mock_http_response
     ):
         with start_transaction(name="google_genai_embeddings"):
             mock_genai_client.models.embed_content(
                 model="text-embedding-004",
-                contents=["Test without metadata"],
+                contents=["Test without statistics", "Another test"],
             )
 
     (event,) = events
     (embed_span,) = event["spans"]
 
-    # But no usage tokens since there's no metadata
+    # No usage tokens since there are no statistics in older versions
+    # This is expected and the integration should handle it gracefully
     assert SPANDATA.GEN_AI_USAGE_INPUT_TOKENS not in embed_span["data"]