From af90c7be0e6157042a343ad1c5776e6ff4337eca Mon Sep 17 00:00:00 2001 From: Alexander Alderman Webb Date: Mon, 7 Sep 2026 16:19:48 +0200 Subject: [PATCH 1/3] test(cohere): Parametrize tests on the streaming trace lifecycle --- tests/integrations/cohere/test_cohere.py | 154 ++++++++++++++++------- 1 file changed, 108 insertions(+), 46 deletions(-) diff --git a/tests/integrations/cohere/test_cohere.py b/tests/integrations/cohere/test_cohere.py index b27b8e5d4e..74506e6fcf 100644 --- a/tests/integrations/cohere/test_cohere.py +++ b/tests/integrations/cohere/test_cohere.py @@ -386,74 +386,136 @@ def test_embed( assert span["data"]["gen_ai.usage.total_tokens"] == 10 -def test_span_origin_chat(sentry_init, capture_events): +@pytest.mark.parametrize("span_streaming", [True, False]) +def test_span_origin_chat(sentry_init, capture_events, capture_items, span_streaming): sentry_init( integrations=[CohereIntegration()], traces_sample_rate=1.0, ) - events = capture_events() - client = Client(api_key="z") - HTTPXClient.request = mock.Mock( - return_value=httpx.Response( - 200, - json={ - "text": "the model response", - "meta": { - "billed_units": { - "output_tokens": 10, - "input_tokens": 20, - } + if span_streaming: + items = capture_events("span") + + client = Client(api_key="z") + HTTPXClient.request = mock.Mock( + return_value=httpx.Response( + 200, + json={ + "text": "the model response", + "meta": { + "billed_units": { + "output_tokens": 10, + "input_tokens": 20, + } + }, }, - }, + ) ) - ) - with start_transaction(name="cohere tx"): - client.chat( - model="some-model", - chat_history=[ChatMessage(role="SYSTEM", message="some context")], - message="hello", - ).text + with start_transaction(name="cohere tx"): + client.chat( + model="some-model", + chat_history=[ChatMessage(role="SYSTEM", message="some context")], + message="hello", + ).text + + sentry_sdk.flush() + (span,) = (item.payload for item in items) + assert span["attributes"]["sentry.origin"] == "auto.ai.cohere" + else: + events = capture_items("span") + + client = Client(api_key="z") + HTTPXClient.request = mock.Mock( + return_value=httpx.Response( + 200, + json={ + "text": "the model response", + "meta": { + "billed_units": { + "output_tokens": 10, + "input_tokens": 20, + } + }, + }, + ) + ) + + with start_transaction(name="cohere tx"): + client.chat( + model="some-model", + chat_history=[ChatMessage(role="SYSTEM", message="some context")], + message="hello", + ).text - (event,) = events + (event,) = events - assert event["contexts"]["trace"]["origin"] == "manual" - assert event["spans"][0]["origin"] == "auto.ai.cohere" + assert event["contexts"]["trace"]["origin"] == "manual" + assert event["spans"][0]["origin"] == "auto.ai.cohere" -def test_span_origin_embed(sentry_init, capture_events): +@pytest.mark.parametrize("span_streaming", [True, False]) +def test_span_origin_embed(sentry_init, capture_events, capture_items, span_streaming): sentry_init( integrations=[CohereIntegration()], traces_sample_rate=1.0, ) - events = capture_events() - client = Client(api_key="z") - HTTPXClient.request = mock.Mock( - return_value=httpx.Response( - 200, - json={ - "response_type": "embeddings_floats", - "id": "1", - "texts": ["hello"], - "embeddings": [[1.0, 2.0, 3.0]], - "meta": { - "billed_units": { - "input_tokens": 10, - } + if span_streaming: + items = capture_items("span") + + client = Client(api_key="z") + HTTPXClient.request = mock.Mock( + return_value=httpx.Response( + 200, + json={ + "response_type": "embeddings_floats", + "id": "1", + "texts": ["hello"], + "embeddings": [[1.0, 2.0, 3.0]], + "meta": { + "billed_units": { + "input_tokens": 10, + } + }, }, - }, + ) ) - ) - with start_transaction(name="cohere tx"): - client.embed(texts=["hello"], model="text-embedding-3-large") + with start_transaction(name="cohere tx"): + client.embed(texts=["hello"], model="text-embedding-3-large") + + sentry_sdk.flush() + (span,) = (item.payload for item in items) + assert span["attributes"]["sentry.origin"] == "auto.ai.cohere" + else: + events = capture_items("span") + + client = Client(api_key="z") + HTTPXClient.request = mock.Mock( + return_value=httpx.Response( + 200, + json={ + "response_type": "embeddings_floats", + "id": "1", + "texts": ["hello"], + "embeddings": [[1.0, 2.0, 3.0]], + "meta": { + "billed_units": { + "input_tokens": 10, + } + }, + }, + ) + ) + + with start_transaction(name="cohere tx"): + client.embed(texts=["hello"], model="text-embedding-3-large") - (event,) = events + (event,) = events - assert event["contexts"]["trace"]["origin"] == "manual" - assert event["spans"][0]["origin"] == "auto.ai.cohere" + assert event["contexts"]["trace"]["origin"] == "manual" + assert event["spans"][0]["origin"] == "auto.ai.cohere" # data_collection config, send_default_pii, include_prompts, expect_inputs, expect_outputs From 9d343f47e05175b2c4e6dc855f63f5807170c8a2 Mon Sep 17 00:00:00 2001 From: Alexander Alderman Webb Date: Mon, 7 Sep 2026 16:23:09 +0200 Subject: [PATCH 2/3] . --- tests/integrations/cohere/test_cohere.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/tests/integrations/cohere/test_cohere.py b/tests/integrations/cohere/test_cohere.py index 74506e6fcf..1bab8bfe95 100644 --- a/tests/integrations/cohere/test_cohere.py +++ b/tests/integrations/cohere/test_cohere.py @@ -391,10 +391,11 @@ def test_span_origin_chat(sentry_init, capture_events, capture_items, span_strea sentry_init( integrations=[CohereIntegration()], traces_sample_rate=1.0, + trace_lifecycle="stream" if span_streaming else "static", ) if span_streaming: - items = capture_events("span") + items = capture_items("span") client = Client(api_key="z") HTTPXClient.request = mock.Mock( @@ -423,7 +424,7 @@ def test_span_origin_chat(sentry_init, capture_events, capture_items, span_strea (span,) = (item.payload for item in items) assert span["attributes"]["sentry.origin"] == "auto.ai.cohere" else: - events = capture_items("span") + events = capture_events() client = Client(api_key="z") HTTPXClient.request = mock.Mock( @@ -459,6 +460,7 @@ def test_span_origin_embed(sentry_init, capture_events, capture_items, span_stre sentry_init( integrations=[CohereIntegration()], traces_sample_rate=1.0, + trace_lifecycle="stream" if span_streaming else "static", ) if span_streaming: @@ -489,7 +491,7 @@ def test_span_origin_embed(sentry_init, capture_events, capture_items, span_stre (span,) = (item.payload for item in items) assert span["attributes"]["sentry.origin"] == "auto.ai.cohere" else: - events = capture_items("span") + events = capture_events() client = Client(api_key="z") HTTPXClient.request = mock.Mock( From 7f0138c344c53dbb9008afd3d9538d7845191259 Mon Sep 17 00:00:00 2001 From: Alexander Alderman Webb Date: Tue, 8 Sep 2026 10:02:11 +0200 Subject: [PATCH 3/3] parametrize test_bad_chat --- tests/integrations/cohere/test_cohere.py | 47 ++++++++++++++++++------ 1 file changed, 35 insertions(+), 12 deletions(-) diff --git a/tests/integrations/cohere/test_cohere.py b/tests/integrations/cohere/test_cohere.py index 1bab8bfe95..db13e39ade 100644 --- a/tests/integrations/cohere/test_cohere.py +++ b/tests/integrations/cohere/test_cohere.py @@ -249,20 +249,43 @@ def test_streaming_chat( assert span["data"]["gen_ai.usage.total_tokens"] == 30 -def test_bad_chat(sentry_init, capture_events): - sentry_init(integrations=[CohereIntegration()], traces_sample_rate=1.0) - events = capture_events() - - client = Client(api_key="z") - HTTPXClient.request = mock.Mock( - side_effect=httpx.HTTPError("API rate limit reached") +@pytest.mark.parametrize("span_streaming", [True, False]) +def test_bad_chat(sentry_init, capture_events, capture_items, span_streaming): + sentry_init( + integrations=[CohereIntegration()], + traces_sample_rate=1.0, + trace_lifecycle="stream" if span_streaming else "static", ) - with pytest.raises(httpx.HTTPError): - client.chat(model="some-model", message="hello") - (event, transaction) = events - assert event["level"] == "error" - assert transaction["contexts"]["trace"]["status"] == "internal_error" + if span_streaming: + items = capture_items("event", "span") + + client = Client(api_key="z") + HTTPXClient.request = mock.Mock( + side_effect=httpx.HTTPError("API rate limit reached") + ) + with pytest.raises(httpx.HTTPError): + client.chat(model="some-model", message="hello") + + (event,) = (item.payload for item in items if item.type == "event") + assert event["level"] == "error" + + sentry_sdk.flush() + (span,) = (item.payload for item in items if item.type == "span") + assert span["status"] == "error" + else: + events = capture_events() + + client = Client(api_key="z") + HTTPXClient.request = mock.Mock( + side_effect=httpx.HTTPError("API rate limit reached") + ) + with pytest.raises(httpx.HTTPError): + client.chat(model="some-model", message="hello") + + (event, transaction) = events + assert event["level"] == "error" + assert transaction["contexts"]["trace"]["status"] == "internal_error" def test_span_status_error(sentry_init, capture_events):