Skip to content

test: Add token span attribute integration tests#312

Open
Pvnn wants to merge 3 commits into
KeyValueSoftwareSystems:developfrom
Pvnn:test/token-span-attribute-tests
Open

test: Add token span attribute integration tests#312
Pvnn wants to merge 3 commits into
KeyValueSoftwareSystems:developfrom
Pvnn:test/token-span-attribute-tests

Conversation

@Pvnn

@Pvnn Pvnn commented Jun 11, 2026

Copy link
Copy Markdown

Overview

Adds integration-style tests for token-to-span attribute assignment across all instrumented LLM providers, covering the full pipeline:

Provider utils → SpanIOProcessor-patched span.set_attribute → final span attributes

What's tested

Provider Happy Path Alt Keys Missing Fields Zero Values Partial Data Special
SpanIOProcessor aliasing
OpenAI cached + reasoning
Groq
Cerebras object vs dict details
LiteLLM reasoning tokens
DSPy
Pydantic AI string cast + dual-write
Google GenAI thoughts summing
OpenAI Streaming chunk accumulation

39 tests, all passing.


Findings during implementation

1. Zero token values silently dropped (all providers) Extraction uses usage.get("prompt_tokens") or usage.get("input_tokens")0 is falsy so it resolves to None before any is not None guard runs. A valid zero completion (e.g. safety-blocked response) leaves no token attributes on the span.

2. Pydantic AI casts token values to strings _safe_set_attribute calls str(value) unconditionally before writing. Token counts land as "100" instead of 100. Downstream tools doing arithmetic on token metrics may need to consider this?


Test approach

A patched_span fixture creates a MagicMock span, wires set_attribute to store into span.attributes, then calls SpanIOProcessor.on_start(span) to apply the aliasing closure — replicating the exact runtime pipeline with no tracer or OTLP overhead.

Comment thread tests/test_token_usage_attributes.py Outdated
Comment thread tests/test_token_usage_attributes.py Outdated

# -- OpenAI streaming --

def test_openai_streaming_accumulation(self, patched_span):

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Either split into two focused tests (one for chunk accumulation, one for finalize → span attributes), or extract a helper that sets up the StreamingWrapper with patches.

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Extracted _make_streaming_wrapper helper and split into focused accumulation and finalization tests.

Comment thread tests/test_token_usage_attributes.py Outdated
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants