Skip to content

Commit e05ea81

Browse files
alexander-alderman-webbalexander-alderman-webb
committed
.
1 parent 4514d4d commit e05ea81

File tree

1 file changed

+75
-99
lines changed

1 file changed

+75
-99
lines changed

‎tests/integrations/openai/test_openai.py‎

Lines changed: 75 additions & 99 deletions
Original file line numberDiff line numberDiff line change
@@ -1594,62 +1594,50 @@ async def test_span_origin_nonstreaming_chat_async(sentry_init, capture_events):
15941594
assert event["spans"][0]["origin"] == "auto.ai.openai"
15951595

15961596

1597-
def test_span_origin_streaming_chat(
1598-
sentry_init, capture_events, get_model_response, server_side_event_chunks
1599-
):
1597+
def test_span_origin_streaming_chat(sentry_init, capture_events):
16001598
sentry_init(
16011599
integrations=[OpenAIIntegration()],
16021600
traces_sample_rate=1.0,
16031601
)
16041602
events = capture_events()
16051603

16061604
client = OpenAI(api_key="z")
1607-
returned_stream = get_model_response(
1608-
server_side_event_chunks(
1609-
[
1610-
ChatCompletionChunk(
1611-
id="1",
1612-
choices=[
1613-
DeltaChoice(
1614-
index=0,
1615-
delta=ChoiceDelta(content="hel"),
1616-
finish_reason=None,
1617-
)
1618-
],
1619-
created=100000,
1620-
model="model-id",
1621-
object="chat.completion.chunk",
1622-
),
1623-
ChatCompletionChunk(
1624-
id="1",
1625-
choices=[
1626-
DeltaChoice(
1627-
index=1,
1628-
delta=ChoiceDelta(content="lo "),
1629-
finish_reason=None,
1630-
)
1631-
],
1632-
created=100000,
1633-
model="model-id",
1634-
object="chat.completion.chunk",
1635-
),
1636-
ChatCompletionChunk(
1637-
id="1",
1638-
choices=[
1639-
DeltaChoice(
1640-
index=2,
1641-
delta=ChoiceDelta(content="world"),
1642-
finish_reason="stop",
1643-
)
1644-
],
1645-
created=100000,
1646-
model="model-id",
1647-
object="chat.completion.chunk",
1648-
),
1605+
returned_stream = Stream(cast_to=None, response=None, client=client)
1606+
returned_stream._iterator = [
1607+
ChatCompletionChunk(
1608+
id="1",
1609+
choices=[
1610+
DeltaChoice(
1611+
index=0, delta=ChoiceDelta(content="hel"), finish_reason=None
1612+
)
16491613
],
1650-
include_event_type=False,
1651-
)
1652-
)
1614+
created=100000,
1615+
model="model-id",
1616+
object="chat.completion.chunk",
1617+
),
1618+
ChatCompletionChunk(
1619+
id="1",
1620+
choices=[
1621+
DeltaChoice(
1622+
index=1, delta=ChoiceDelta(content="lo "), finish_reason=None
1623+
)
1624+
],
1625+
created=100000,
1626+
model="model-id",
1627+
object="chat.completion.chunk",
1628+
),
1629+
ChatCompletionChunk(
1630+
id="1",
1631+
choices=[
1632+
DeltaChoice(
1633+
index=2, delta=ChoiceDelta(content="world"), finish_reason="stop"
1634+
)
1635+
],
1636+
created=100000,
1637+
model="model-id",
1638+
object="chat.completion.chunk",
1639+
),
1640+
]
16531641

16541642
client.chat.completions._post = mock.Mock(return_value=returned_stream)
16551643
with start_transaction(name="openai tx"):
@@ -1667,11 +1655,7 @@ def test_span_origin_streaming_chat(
16671655

16681656
@pytest.mark.asyncio
16691657
async def test_span_origin_streaming_chat_async(
1670-
sentry_init,
1671-
capture_events,
1672-
get_model_response,
1673-
async_iterator,
1674-
server_side_event_chunks,
1658+
sentry_init, capture_events, async_iterator
16751659
):
16761660
sentry_init(
16771661
integrations=[OpenAIIntegration()],
@@ -1680,53 +1664,45 @@ async def test_span_origin_streaming_chat_async(
16801664
events = capture_events()
16811665

16821666
client = AsyncOpenAI(api_key="z")
1683-
returned_stream = get_model_response(
1684-
async_iterator(
1685-
server_side_event_chunks(
1686-
[
1687-
ChatCompletionChunk(
1688-
id="1",
1689-
choices=[
1690-
DeltaChoice(
1691-
index=0,
1692-
delta=ChoiceDelta(content="hel"),
1693-
finish_reason=None,
1694-
)
1695-
],
1696-
created=100000,
1697-
model="model-id",
1698-
object="chat.completion.chunk",
1699-
),
1700-
ChatCompletionChunk(
1701-
id="1",
1702-
choices=[
1703-
DeltaChoice(
1704-
index=1,
1705-
delta=ChoiceDelta(content="lo "),
1706-
finish_reason=None,
1707-
)
1708-
],
1709-
created=100000,
1710-
model="model-id",
1711-
object="chat.completion.chunk",
1712-
),
1713-
ChatCompletionChunk(
1714-
id="1",
1715-
choices=[
1716-
DeltaChoice(
1717-
index=2,
1718-
delta=ChoiceDelta(content="world"),
1719-
finish_reason="stop",
1720-
)
1721-
],
1722-
created=100000,
1723-
model="model-id",
1724-
object="chat.completion.chunk",
1725-
),
1667+
returned_stream = AsyncStream(cast_to=None, response=None, client=client)
1668+
returned_stream._iterator = async_iterator(
1669+
[
1670+
ChatCompletionChunk(
1671+
id="1",
1672+
choices=[
1673+
DeltaChoice(
1674+
index=0, delta=ChoiceDelta(content="hel"), finish_reason=None
1675+
)
17261676
],
1727-
include_event_type=False,
1728-
)
1729-
)
1677+
created=100000,
1678+
model="model-id",
1679+
object="chat.completion.chunk",
1680+
),
1681+
ChatCompletionChunk(
1682+
id="1",
1683+
choices=[
1684+
DeltaChoice(
1685+
index=1, delta=ChoiceDelta(content="lo "), finish_reason=None
1686+
)
1687+
],
1688+
created=100000,
1689+
model="model-id",
1690+
object="chat.completion.chunk",
1691+
),
1692+
ChatCompletionChunk(
1693+
id="1",
1694+
choices=[
1695+
DeltaChoice(
1696+
index=2,
1697+
delta=ChoiceDelta(content="world"),
1698+
finish_reason="stop",
1699+
)
1700+
],
1701+
created=100000,
1702+
model="model-id",
1703+
object="chat.completion.chunk",
1704+
),
1705+
]
17301706
)
17311707

17321708
client.chat.completions._post = AsyncMock(return_value=returned_stream)

0 commit comments

Comments
 (0)