Skip to content

Commit b968b8e

Browse files
shellmayrshellmayr
committed
add tests
1 parent fb14ce0 commit b968b8e

4 files changed

Lines changed: 833 additions & 1 deletion

File tree

‎tests/integrations/anthropic/test_anthropic.py‎

Lines changed: 267 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1446,3 +1446,270 @@ def test_system_prompt_with_complex_structure(sentry_init, capture_events):
14461446
assert stored_messages[0]["content"][1]["text"] == "Be concise and clear."
14471447
assert stored_messages[1]["role"] == "user"
14481448
assert stored_messages[1]["content"] == "Hello"
1449+
1450+
1451+
def test_nonstreaming_create_message_with_cache_tokens(sentry_init, capture_events):
1452+
"""Test that cache tokens (read and write) are captured from Anthropic responses."""
1453+
1454+
# Create a mock usage object with cache tokens
1455+
class MockUsageWithCache:
1456+
def __init__(self):
1457+
self.input_tokens = 100
1458+
self.output_tokens = 50
1459+
self.cache_read_input_tokens = 25
1460+
self.cache_write_input_tokens = 30
1461+
1462+
# Create a mock message with cache tokens
1463+
example_message_with_cache = Message(
1464+
id="id",
1465+
model="model",
1466+
role="assistant",
1467+
content=[TextBlock(type="text", text="Hi, I'm Claude.")],
1468+
type="message",
1469+
usage=Usage(input_tokens=100, output_tokens=50),
1470+
)
1471+
# Add cache token attributes to the usage object
1472+
example_message_with_cache.usage.cache_read_input_tokens = 25
1473+
example_message_with_cache.usage.cache_write_input_tokens = 30
1474+
1475+
sentry_init(
1476+
integrations=[AnthropicIntegration()],
1477+
traces_sample_rate=1.0,
1478+
)
1479+
events = capture_events()
1480+
client = Anthropic(api_key="z")
1481+
client.messages._post = mock.Mock(return_value=example_message_with_cache)
1482+
1483+
messages = [{"role": "user", "content": "Hello, Claude"}]
1484+
1485+
with start_transaction(name="anthropic"):
1486+
response = client.messages.create(
1487+
max_tokens=1024, messages=messages, model="model"
1488+
)
1489+
1490+
assert response == example_message_with_cache
1491+
1492+
assert len(events) == 1
1493+
(event,) = events
1494+
1495+
assert len(event["spans"]) == 1
1496+
(span,) = event["spans"]
1497+
1498+
assert span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 100
1499+
assert span["data"][SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 50
1500+
assert span["data"][SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS] == 150
1501+
assert span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED] == 25
1502+
assert span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHE_WRITE] == 30
1503+
1504+
1505+
@pytest.mark.asyncio
1506+
async def test_nonstreaming_create_message_with_cache_tokens_async(
1507+
sentry_init, capture_events
1508+
):
1509+
"""Test that cache tokens are captured from async Anthropic responses."""
1510+
example_message_with_cache = Message(
1511+
id="id",
1512+
model="model",
1513+
role="assistant",
1514+
content=[TextBlock(type="text", text="Hi, I'm Claude.")],
1515+
type="message",
1516+
usage=Usage(input_tokens=100, output_tokens=50),
1517+
)
1518+
example_message_with_cache.usage.cache_read_input_tokens = 25
1519+
example_message_with_cache.usage.cache_write_input_tokens = 30
1520+
1521+
sentry_init(
1522+
integrations=[AnthropicIntegration()],
1523+
traces_sample_rate=1.0,
1524+
)
1525+
events = capture_events()
1526+
client = AsyncAnthropic(api_key="z")
1527+
client.messages._post = AsyncMock(return_value=example_message_with_cache)
1528+
1529+
messages = [{"role": "user", "content": "Hello, Claude"}]
1530+
1531+
with start_transaction(name="anthropic"):
1532+
response = await client.messages.create(
1533+
max_tokens=1024, messages=messages, model="model"
1534+
)
1535+
1536+
assert response == example_message_with_cache
1537+
1538+
assert len(events) == 1
1539+
(event,) = events
1540+
1541+
(span,) = event["spans"]
1542+
1543+
assert span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 100
1544+
assert span["data"][SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 50
1545+
assert span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED] == 25
1546+
assert span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHE_WRITE] == 30
1547+
1548+
1549+
def test_streaming_create_message_with_cache_tokens(sentry_init, capture_events):
1550+
"""Test that cache tokens are captured from streaming Anthropic responses."""
1551+
# Create a mock usage object with cache tokens for the message start event
1552+
example_message_with_cache = Message(
1553+
id="id",
1554+
model="model",
1555+
role="assistant",
1556+
content=[TextBlock(type="text", text="Hi, I'm Claude.")],
1557+
type="message",
1558+
usage=Usage(input_tokens=100, output_tokens=20),
1559+
)
1560+
example_message_with_cache.usage.cache_read_input_tokens = 25
1561+
example_message_with_cache.usage.cache_write_input_tokens = 30
1562+
1563+
client = Anthropic(api_key="z")
1564+
returned_stream = Stream(cast_to=None, response=None, client=client)
1565+
returned_stream._iterator = [
1566+
MessageStartEvent(
1567+
message=example_message_with_cache,
1568+
type="message_start",
1569+
),
1570+
ContentBlockStartEvent(
1571+
type="content_block_start",
1572+
index=0,
1573+
content_block=TextBlock(type="text", text=""),
1574+
),
1575+
ContentBlockDeltaEvent(
1576+
delta=TextDelta(text="Hi!", type="text_delta"),
1577+
index=0,
1578+
type="content_block_delta",
1579+
),
1580+
ContentBlockStopEvent(type="content_block_stop", index=0),
1581+
MessageDeltaEvent(
1582+
delta=Delta(),
1583+
usage=MessageDeltaUsage(output_tokens=10),
1584+
type="message_delta",
1585+
),
1586+
]
1587+
1588+
sentry_init(
1589+
integrations=[AnthropicIntegration()],
1590+
traces_sample_rate=1.0,
1591+
)
1592+
events = capture_events()
1593+
client.messages._post = mock.Mock(return_value=returned_stream)
1594+
1595+
messages = [{"role": "user", "content": "Hello, Claude"}]
1596+
1597+
with start_transaction(name="anthropic"):
1598+
message = client.messages.create(
1599+
max_tokens=1024, messages=messages, model="model", stream=True
1600+
)
1601+
for _ in message:
1602+
pass
1603+
1604+
assert len(events) == 1
1605+
(event,) = events
1606+
1607+
(span,) = event["spans"]
1608+
1609+
# For streaming, tokens accumulate from message_start and message_delta
1610+
assert span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 100
1611+
# Output tokens: 20 from message_start + 10 from message_delta = 30
1612+
assert span["data"][SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 30
1613+
assert span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED] == 25
1614+
assert span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHE_WRITE] == 30
1615+
1616+
1617+
@pytest.mark.asyncio
1618+
async def test_streaming_create_message_with_cache_tokens_async(
1619+
sentry_init, capture_events
1620+
):
1621+
"""Test that cache tokens are captured from async streaming Anthropic responses."""
1622+
example_message_with_cache = Message(
1623+
id="id",
1624+
model="model",
1625+
role="assistant",
1626+
content=[TextBlock(type="text", text="Hi, I'm Claude.")],
1627+
type="message",
1628+
usage=Usage(input_tokens=100, output_tokens=20),
1629+
)
1630+
example_message_with_cache.usage.cache_read_input_tokens = 25
1631+
example_message_with_cache.usage.cache_write_input_tokens = 30
1632+
1633+
client = AsyncAnthropic(api_key="z")
1634+
returned_stream = AsyncStream(cast_to=None, response=None, client=client)
1635+
returned_stream._iterator = async_iterator(
1636+
[
1637+
MessageStartEvent(
1638+
message=example_message_with_cache,
1639+
type="message_start",
1640+
),
1641+
ContentBlockStartEvent(
1642+
type="content_block_start",
1643+
index=0,
1644+
content_block=TextBlock(type="text", text=""),
1645+
),
1646+
ContentBlockDeltaEvent(
1647+
delta=TextDelta(text="Hi!", type="text_delta"),
1648+
index=0,
1649+
type="content_block_delta",
1650+
),
1651+
ContentBlockStopEvent(type="content_block_stop", index=0),
1652+
MessageDeltaEvent(
1653+
delta=Delta(),
1654+
usage=MessageDeltaUsage(output_tokens=10),
1655+
type="message_delta",
1656+
),
1657+
]
1658+
)
1659+
1660+
sentry_init(
1661+
integrations=[AnthropicIntegration()],
1662+
traces_sample_rate=1.0,
1663+
)
1664+
events = capture_events()
1665+
client.messages._post = AsyncMock(return_value=returned_stream)
1666+
1667+
messages = [{"role": "user", "content": "Hello, Claude"}]
1668+
1669+
with start_transaction(name="anthropic"):
1670+
message = await client.messages.create(
1671+
max_tokens=1024, messages=messages, model="model", stream=True
1672+
)
1673+
async for _ in message:
1674+
pass
1675+
1676+
assert len(events) == 1
1677+
(event,) = events
1678+
1679+
(span,) = event["spans"]
1680+
1681+
assert span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 100
1682+
assert span["data"][SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 30
1683+
assert span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED] == 25
1684+
assert span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHE_WRITE] == 30
1685+
1686+
1687+
def test_nonstreaming_without_cache_tokens(sentry_init, capture_events):
1688+
"""Test that missing or zero cache tokens don't cause errors (backwards compatibility)."""
1689+
sentry_init(
1690+
integrations=[AnthropicIntegration()],
1691+
traces_sample_rate=1.0,
1692+
)
1693+
events = capture_events()
1694+
client = Anthropic(api_key="z")
1695+
# Using EXAMPLE_MESSAGE which doesn't have explicit cache tokens set
1696+
client.messages._post = mock.Mock(return_value=EXAMPLE_MESSAGE)
1697+
1698+
messages = [{"role": "user", "content": "Hello, Claude"}]
1699+
1700+
with start_transaction(name="anthropic"):
1701+
client.messages.create(max_tokens=1024, messages=messages, model="model")
1702+
1703+
assert len(events) == 1
1704+
(event,) = events
1705+
1706+
(span,) = event["spans"]
1707+
1708+
assert span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 10
1709+
assert span["data"][SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 20
1710+
# Cache tokens may be present with 0 values (newer SDK versions)
1711+
# or not present at all (older SDK versions) - both are acceptable
1712+
cached = span["data"].get(SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED)
1713+
cache_write = span["data"].get(SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHE_WRITE)
1714+
assert cached is None or cached == 0
1715+
assert cache_write is None or cache_write == 0

0 commit comments

Comments
 (0)