add more tests around collecting ai data

Author: kwnath

tests/integrations/anthropic/test_anthropic.py

@@ -10,7 +10,7 @@ async def __call__(self, *args, **kwargs):
 
 
 import pytest
-from anthropic import AsyncAnthropic, Anthropic, AnthropicError, AsyncStream, Stream
+from anthropic import Anthropic, AnthropicError, AsyncAnthropic, AsyncStream, Stream
 from anthropic.types import MessageDeltaUsage, TextDelta, Usage
 from anthropic.types.content_block_delta_event import ContentBlockDeltaEvent
 from anthropic.types.content_block_start_event import ContentBlockStartEvent
@@ -19,6 +19,7 @@ async def __call__(self, *args, **kwargs):
 from anthropic.types.message_delta_event import MessageDeltaEvent
 from anthropic.types.message_start_event import MessageStartEvent
 
+from sentry_sdk.integrations.anthropic import _add_ai_data_to_span, _collect_ai_data
 from sentry_sdk.utils import package_version
 
 try:
@@ -757,3 +758,59 @@ async def test_span_origin_async(sentry_init, capture_events):
 
     assert event["contexts"]["trace"]["origin"] == "manual"
     assert event["spans"][0]["origin"] == "auto.ai.anthropic"
+
+
+@pytest.mark.skipif(
+    ANTHROPIC_VERSION < (0, 27),
+    reason="Versions <0.27.0 do not include InputJSONDelta.",
+)
+def test_collect_ai_data_with_input_json_delta():
+    event = ContentBlockDeltaEvent(
+        delta=InputJSONDelta(partial_json="test", type="input_json_delta"),
+        index=0,
+        type="content_block_delta",
+    )
+
+    input_tokens = 10
+    output_tokens = 20
+    content_blocks = []
+
+    new_input_tokens, new_output_tokens, new_content_blocks = _collect_ai_data(
+        event, input_tokens, output_tokens, content_blocks
+    )
+
+    assert new_input_tokens == input_tokens
+    assert new_output_tokens == output_tokens
+    assert new_content_blocks == ["test"]
+
+
+@pytest.mark.skipif(
+    ANTHROPIC_VERSION < (0, 27),
+    reason="Versions <0.27.0 do not include InputJSONDelta.",
+)
+def test_add_ai_data_to_span_with_input_json_delta(sentry_init):
+    sentry_init(
+        integrations=[AnthropicIntegration(include_prompts=True)],
+        traces_sample_rate=1.0,
+        send_default_pii=True,
+    )
+
+    with start_transaction(name="test") as transaction:
+        span = transaction.start_span()
+        integration = AnthropicIntegration()
+
+        _add_ai_data_to_span(
+            span,
+            integration,
+            input_tokens=10,
+            output_tokens=20,
+            content_blocks=["{'test': 'data'}", "'more': 'json'"],
+        )
+
+        assert span.get_data(SPANDATA.AI_RESPONSES) == [
+            {"type": "text", "text": "{'test': 'data''more': 'json'}"}
+        ]
+        assert span.get_data("ai.streaming") is True
+        assert span.get_measurement("ai_prompt_tokens_used")["value"] == 10
+        assert span.get_measurement("ai_completion_tokens_used")["value"] == 20
+        assert span.get_measurement("ai_total_tokens_used")["value"] == 30