def test_get_token_usage_for_event():
    """Test that we get the single matching usage record (if any) based on the event's model_response.id."""
    metrics = Metrics(model_name='test-model')
    usage_record = TokenUsage(
        model='test-model',
        prompt_tokens=10,
        completion_tokens=5,
        cache_read_tokens=2,
        cache_write_tokens=1,
        response_id='test-response-id',
    )
    metrics.add_token_usage(
        prompt_tokens=usage_record.prompt_tokens,
        completion_tokens=usage_record.completion_tokens,
        cache_read_tokens=usage_record.cache_read_tokens,
        cache_write_tokens=usage_record.cache_write_tokens,
        context_window=1000,
        response_id=usage_record.response_id,
    )

    # Create an event referencing that response_id
    event = Event()
    mock_tool_call_metadata = ToolCallMetadata(
        tool_call_id='test-tool-call',
        function_name='fake_function',
        model_response={'id': 'test-response-id'},
        total_calls_in_response=1,
    )
    event._tool_call_metadata = (
        mock_tool_call_metadata  # normally you'd do event.tool_call_metadata = ...
    )

    # We should find that usage record
    found = get_token_usage_for_event(event, metrics)
    assert found is not None
    assert found.prompt_tokens == 10
    assert found.response_id == 'test-response-id'

    # If we change the event's response ID, we won't find anything
    mock_tool_call_metadata.model_response.id = 'some-other-id'
    found2 = get_token_usage_for_event(event, metrics)
    assert found2 is None

    # If the event has no tool_call_metadata, also returns None
    event._tool_call_metadata = None
    found3 = get_token_usage_for_event(event, metrics)
    assert found3 is None
Test that we get the single matching usage record (if any) based on the event's model_response.id.
tests/unit/test_message_utils.py