Praneel Panchigar
7056c1c687
Feature/automated cost logging csv
* feat: Add usage and cost tracking with CSV export
- Add usage_tracking.py module for thread-safe token and cost tracking
- Extract actual tokens from Bedrock API responses (replaces word count estimation)
- Integrate usage tracking into llm_utils.py for all LLM invocations
- Add CSV export functionality (per-file/per-model and batch summary)
- Integrate CSV export into main.py for local runs
- Update .gitignore to exclude PRD
Phase 1: Data collection and local CSV export implemented
- Tracks all LLM calls including utility functions (date_fix, derive_term_date)
- Calculates costs per model using centralized cost constants
- Generates two CSV files: USAGE.csv and USAGE-SUMMARY.csv
- Exports to same directory as final results when write_to_s3=False
* - Removed some temporary logging code
- Removed unused imports (uuid, datetime) from llm_utils.py
- Added cache token columns (cache_creation_tokens, cache_read_tokens) to CSV exports
- Added tracking for cache warming
- Updated usage_tracking.py to track cache tokens separately in data structure
All token tracking functionality remains intact.
* Split input tokens into fresh and cache in summary section
- Add total_fresh_input_tokens and total_cache_tokens to GLOBAL_USAGE
- Update summary CSV to include split token columns for better tractability
- Add averages for fresh_input_tokens and cache_tokens
- Maintain total_input_tokens = fresh_input_tokens + cache_tokens relationship
- All calculations verified and match per-file totals correctly
* upload files functionality
* Add comprehensive unit tests for usage_tracking module
- Implement Phase 1-4 tests covering all 10 functions in usage_tracking.py
- Add 98 test cases with parametrized tests for comprehensive coverage
- Fix extract_usage_from_response to handle None input gracefully
- Add python-dotenv dependency to pyproject.toml
- Tests cover: core functions, data access, export functions, and utilities
- All tests passing (98/98)
* Add S3 upload tests for main.py and fix f-string syntax error
- Add tests/test_investment_main_s3.py with 3 test cases:
* test_main_writes_final_and_error_to_s3: validates both final and error DataFrames uploaded when processing has errors
* test_main_writes_only_final_when_no_errors: validates only final DataFrame uploaded when all files succeed
* test_main_writes_usage_when_present: validates usage tracking files uploaded when usage data exists
- Tests use monkeypatching and module stubs to isolate S3 write logic without external dependencies
- Fix syntax error in src/codes/code_funcs.py: nested f-string quotes (level_dict['level_suffix'])
- All tests passing
* Add comprehensive test for usage_tracking.get_usage_dataframes S3 uploads
- Add test_usage_dataframes_written_to_s3 to validate per_file_usage_df and batch_summary_df
- Test verifies both dataframes from usage_tracking.get_usage_dataframes() are written to S3
- Validates 'usage' suffix writes per_file_usage_df with correct schema (file_name, tokens, cost)
- Validates 'usage_summary' suffix writes batch_summary_df with correct schema (totals, averages, region_mode)
- Uses pd.testing.assert_frame_equal to ensure exact dataframes are uploaded
- Fix module caching issue between tests by clearing sys.modules cache
- All 4 tests passing
* Add cleanup fixture to prevent test pollution
- Add autouse pytest fixture to clean up stubbed modules after each test
- Remove manual cache clearing from individual tests (now handled by fixture)
- Prevents our module stubs from affecting other test files in CI/CD pipeline
- Fixes Bitbucket pipeline failures where other tests couldn't find llm_utils attributes
* Refactor tests: add comprehensive io_utils coverage; disable usage_tracking suite
- Added 22 focused tests for write_local and write_s3 plus existing IO behaviors
- Introduced FakeS3Client to avoid boto3 network/credential dependency
- Added preserve_config fixture (no monkeypatch) to isolate config side effects
- Marked test_usage_tracking.py skipped per new consolidation approach
- Verified 26 tests (io_utils + investment_main) all pass locally without AWS exceptions
* Fix S3_CLIENT mocking: use mocker.patch instead of direct assignment
- Changed preserve_config fixture to NOT save/restore S3_CLIENT
- Updated all 5 write_s3 tests to use mocker.patch('src.config.S3_CLIENT', fake)
- This prevents pollution of config.S3_CLIENT across test modules
- Fixes NoCredentialsError in other tests that were importing config after our tests set FakeS3Client
- All 22 io_utils tests pass + 4 investment_main_s3 tests pass
- test_llm_utils::test_invoke_claude_local_mode now passes
* Add tests for write_s3 function
* Resolve merge conflict: keep mocker-based write_s3 tests
* Remove rogue skip mark from test_usage_tracking.py
The skip mark was incorrectly added by remote branch claiming tests were
migrated to io_utils. However, test_io_utils.py only tests io_utils functions
(read/write operations), not usage_tracking functions.
Restore the comprehensive usage_tracking tests from commit 26a556f0 which
includes 98 test cases covering all 10 functions in usage_tracking.py.
* Fix test pollution: add config cleanup fixture and restore RUN_MODE
- Add autouse fixture in test_io_utils.py to restore config values after each test
- Fix test_llm_utils.py to restore RUN_MODE after test_invoke_claude_local_mode
- Prevents config patches from leaking into subsequent tests causing NoCredentialsError
* Remove redundant test_investment_main_s3.py
- File was testing same write_s3 functionality already covered in test_io_utils.py
- Was causing module pollution by stubbing src.utils.llm_utils
- Caused test failures in other test files due to module cache issues
- Integration testing was minimal and mocked write_s3 anyway
* Merged main into feature/automated-cost-logging-csv
Approved-by: Karan Desai
Approved-by: Katon Minhas
1601 lines
62 KiB
Python
1601 lines
62 KiB
Python
"""Unit tests for usage_tracking module.
|
|
|
|
Phase 1: Core Functions
|
|
- extract_usage_from_response()
|
|
- get_cost_per_token()
|
|
- track_usage()
|
|
|
|
Phase 2: Data Access Functions
|
|
- get_usage_data()
|
|
- reset_usage_data()
|
|
|
|
Phase 3: Export Functions
|
|
- get_usage_dataframes()
|
|
- _generate_per_file_per_model_df()
|
|
- _generate_batch_summary_df()
|
|
- export_usage_cost_csvs_local()
|
|
|
|
Phase 4: Utility Functions
|
|
- get_model_name()
|
|
"""
|
|
|
|
import logging
|
|
import threading
|
|
import time
|
|
|
|
import pandas as pd
|
|
import pytest
|
|
|
|
from src.utils import usage_tracking
|
|
|
|
|
|
class TestExtractUsageFromResponse:
|
|
"""Tests for extract_usage_from_response() function."""
|
|
|
|
@pytest.mark.parametrize(
|
|
"response_body, expected",
|
|
[
|
|
# Valid response with all token types
|
|
(
|
|
{
|
|
"usage": {
|
|
"input_tokens": 100,
|
|
"output_tokens": 50,
|
|
"cache_creation_input_tokens": 20,
|
|
"cache_read_input_tokens": 30,
|
|
}
|
|
},
|
|
(100, 50, 20, 30),
|
|
),
|
|
# Response with only input/output (no cache tokens)
|
|
(
|
|
{
|
|
"usage": {
|
|
"input_tokens": 100,
|
|
"output_tokens": 50,
|
|
}
|
|
},
|
|
(100, 50, 0, 0),
|
|
),
|
|
# Response with cache tokens only
|
|
(
|
|
{
|
|
"usage": {
|
|
"cache_creation_input_tokens": 20,
|
|
"cache_read_input_tokens": 30,
|
|
}
|
|
},
|
|
(0, 0, 20, 30),
|
|
),
|
|
# Response with missing optional fields
|
|
(
|
|
{
|
|
"usage": {
|
|
"input_tokens": 100,
|
|
}
|
|
},
|
|
(100, 0, 0, 0),
|
|
),
|
|
# Response with zero values
|
|
(
|
|
{
|
|
"usage": {
|
|
"input_tokens": 0,
|
|
"output_tokens": 0,
|
|
"cache_creation_input_tokens": 0,
|
|
"cache_read_input_tokens": 0,
|
|
}
|
|
},
|
|
(0, 0, 0, 0),
|
|
),
|
|
# Response with string numbers (should convert to int)
|
|
(
|
|
{
|
|
"usage": {
|
|
"input_tokens": "100",
|
|
"output_tokens": "50",
|
|
"cache_creation_input_tokens": "20",
|
|
"cache_read_input_tokens": "30",
|
|
}
|
|
},
|
|
(100, 50, 20, 30),
|
|
),
|
|
],
|
|
)
|
|
def test_extract_usage_from_response_valid(self, response_body, expected):
|
|
"""Test extraction with valid response bodies."""
|
|
result = usage_tracking.extract_usage_from_response(response_body)
|
|
assert result == expected
|
|
# Verify all values are integers
|
|
assert all(isinstance(x, int) for x in result)
|
|
|
|
@pytest.mark.parametrize(
|
|
"response_body, expected",
|
|
[
|
|
# Missing "usage" key
|
|
({}, (0, 0, 0, 0)),
|
|
# Empty "usage" dict
|
|
({"usage": {}}, (0, 0, 0, 0)),
|
|
# None response
|
|
(None, (0, 0, 0, 0)),
|
|
],
|
|
)
|
|
def test_extract_usage_from_response_missing_usage(self, response_body, expected, caplog):
|
|
"""Test extraction when usage data is missing."""
|
|
with caplog.at_level(logging.DEBUG):
|
|
result = usage_tracking.extract_usage_from_response(response_body)
|
|
assert result == expected
|
|
if response_body is None or "usage" not in (response_body or {}):
|
|
assert "No usage data in API response" in caplog.text
|
|
|
|
def test_extract_usage_from_response_nested_structure(self):
|
|
"""Test that function handles nested structure correctly."""
|
|
response_body = {
|
|
"other_data": "value",
|
|
"usage": {
|
|
"input_tokens": 100,
|
|
"output_tokens": 50,
|
|
"cache_creation_input_tokens": 20,
|
|
"cache_read_input_tokens": 30,
|
|
},
|
|
"more_data": {"nested": "value"},
|
|
}
|
|
result = usage_tracking.extract_usage_from_response(response_body)
|
|
assert result == (100, 50, 20, 30)
|
|
|
|
|
|
class TestGetCostPerToken:
|
|
"""Tests for get_cost_per_token() function."""
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def reset_before_test(self):
|
|
"""Reset usage data before each test."""
|
|
usage_tracking.reset_usage_data()
|
|
yield
|
|
usage_tracking.reset_usage_data()
|
|
|
|
@pytest.mark.parametrize(
|
|
"model_id, expected_input, expected_output, expected_cache_read",
|
|
[
|
|
# Claude 2
|
|
(
|
|
"anthropic.claude-instant-v1",
|
|
0.008 / 1000.0, # $0.008 per 1k tokens
|
|
0.024 / 1000.0, # $0.024 per 1k tokens
|
|
0.0008 / 1000.0, # $0.0008 per 1k tokens
|
|
),
|
|
# Claude 3 Haiku
|
|
(
|
|
"anthropic.claude-3-haiku-20240307-v1:0",
|
|
0.00025 / 1000.0,
|
|
0.00125 / 1000.0,
|
|
0.000025 / 1000.0,
|
|
),
|
|
# Claude 3.5 Sonnet
|
|
(
|
|
"anthropic.claude-3-5-sonnet-20240620-v1:0",
|
|
0.003 / 1000.0,
|
|
0.015 / 1000.0,
|
|
0.0003 / 1000.0,
|
|
),
|
|
(
|
|
"anthropic.claude-3-5-sonnet-20241022-v2:0",
|
|
0.003 / 1000.0,
|
|
0.015 / 1000.0,
|
|
0.0003 / 1000.0,
|
|
),
|
|
# Claude 3.7 Sonnet
|
|
(
|
|
"us.anthropic.claude-3-7-sonnet-20250219-v1:0",
|
|
0.003 / 1000.0,
|
|
0.015 / 1000.0,
|
|
0.0003 / 1000.0,
|
|
),
|
|
# Claude 4 Sonnet
|
|
(
|
|
"us.anthropic.claude-sonnet-4-20250514-v1:0",
|
|
0.003 / 1000.0,
|
|
0.015 / 1000.0,
|
|
0.0003 / 1000.0,
|
|
),
|
|
# Claude 4.5 Sonnet
|
|
(
|
|
"us.anthropic.claude-sonnet-4-5-20250929-v1:0",
|
|
0.003 / 1000.0,
|
|
0.015 / 1000.0,
|
|
0.0003 / 1000.0,
|
|
),
|
|
],
|
|
)
|
|
def test_get_cost_per_token_valid_models(
|
|
self, model_id, expected_input, expected_output, expected_cache_read, mocker
|
|
):
|
|
"""Test cost calculation for valid model IDs."""
|
|
# Mock resolve_model_id to return the model_id as-is
|
|
mocker.patch("src.utils.usage_tracking.config.resolve_model_id", return_value=model_id)
|
|
|
|
input_cost, output_cost, cache_read_cost = usage_tracking.get_cost_per_token(model_id)
|
|
|
|
assert abs(input_cost - expected_input) < 1e-10
|
|
assert abs(output_cost - expected_output) < 1e-10
|
|
assert abs(cache_read_cost - expected_cache_read) < 1e-10
|
|
|
|
def test_get_cost_per_token_unknown_model(self, mocker, caplog):
|
|
"""Test cost calculation for unknown model (should use default)."""
|
|
unknown_model = "unknown.model.id"
|
|
# Mock resolve_model_id to return the unknown model
|
|
mocker.patch(
|
|
"src.utils.usage_tracking.config.resolve_model_id", return_value=unknown_model
|
|
)
|
|
|
|
with caplog.at_level(logging.WARNING):
|
|
input_cost, output_cost, cache_read_cost = usage_tracking.get_cost_per_token(
|
|
unknown_model
|
|
)
|
|
|
|
# Should use default pricing (Claude 3.5 Sonnet)
|
|
expected_input = 0.003 / 1000.0
|
|
expected_output = 0.015 / 1000.0
|
|
expected_cache_read = 0.0003 / 1000.0
|
|
|
|
assert abs(input_cost - expected_input) < 1e-10
|
|
assert abs(output_cost - expected_output) < 1e-10
|
|
assert abs(cache_read_cost - expected_cache_read) < 1e-10
|
|
assert "Unknown model" in caplog.text
|
|
assert "using default pricing" in caplog.text
|
|
|
|
def test_get_cost_per_token_model_alias(self, mocker):
|
|
"""Test that model aliases are resolved correctly."""
|
|
alias = "sonnet_37"
|
|
resolved_id = "us.anthropic.claude-3-7-sonnet-20250219-v1:0"
|
|
# Mock resolve_model_id to resolve the alias
|
|
mocker.patch(
|
|
"src.utils.usage_tracking.config.resolve_model_id", return_value=resolved_id
|
|
)
|
|
|
|
input_cost, output_cost, cache_read_cost = usage_tracking.get_cost_per_token(alias)
|
|
|
|
# Should use Claude 3.7 Sonnet pricing
|
|
expected_input = 0.003 / 1000.0
|
|
expected_output = 0.015 / 1000.0
|
|
expected_cache_read = 0.0003 / 1000.0
|
|
|
|
assert abs(input_cost - expected_input) < 1e-10
|
|
assert abs(output_cost - expected_output) < 1e-10
|
|
assert abs(cache_read_cost - expected_cache_read) < 1e-10
|
|
|
|
def test_get_cost_per_token_conversion(self, mocker):
|
|
"""Test that per-1k costs are correctly converted to per-token."""
|
|
model_id = "anthropic.claude-3-5-sonnet-20240620-v1:0"
|
|
mocker.patch("src.utils.usage_tracking.config.resolve_model_id", return_value=model_id)
|
|
|
|
input_cost, output_cost, cache_read_cost = usage_tracking.get_cost_per_token(model_id)
|
|
|
|
# Verify conversion: per-1k / 1000 = per-token
|
|
assert input_cost == 0.003 / 1000.0
|
|
assert output_cost == 0.015 / 1000.0
|
|
assert cache_read_cost == 0.0003 / 1000.0
|
|
|
|
# Verify all costs are floats
|
|
assert isinstance(input_cost, float)
|
|
assert isinstance(output_cost, float)
|
|
assert isinstance(cache_read_cost, float)
|
|
|
|
|
|
class TestTrackUsage:
|
|
"""Tests for track_usage() function."""
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def reset_before_test(self):
|
|
"""Reset usage data before each test."""
|
|
usage_tracking.reset_usage_data()
|
|
yield
|
|
usage_tracking.reset_usage_data()
|
|
|
|
@pytest.mark.parametrize(
|
|
"filename, model_id, input_tokens, output_tokens, cache_creation, cache_read",
|
|
[
|
|
# All token types present
|
|
("file1.txt", "anthropic.claude-3-5-sonnet-20240620-v1:0", 100, 50, 20, 30),
|
|
# No cache tokens
|
|
("file2.txt", "anthropic.claude-3-5-sonnet-20240620-v1:0", 100, 50, 0, 0),
|
|
# Only cache tokens
|
|
("file3.txt", "anthropic.claude-3-5-sonnet-20240620-v1:0", 0, 0, 20, 30),
|
|
# Zero tokens
|
|
("file4.txt", "anthropic.claude-3-5-sonnet-20240620-v1:0", 0, 0, 0, 0),
|
|
],
|
|
)
|
|
def test_track_usage_single_call(
|
|
self,
|
|
filename,
|
|
model_id,
|
|
input_tokens,
|
|
output_tokens,
|
|
cache_creation,
|
|
cache_read,
|
|
mocker,
|
|
):
|
|
"""Test tracking a single API call."""
|
|
# Mock get_cost_per_token to return known costs
|
|
input_cost = 0.003 / 1000.0
|
|
output_cost = 0.015 / 1000.0
|
|
cache_read_cost = 0.0003 / 1000.0
|
|
mocker.patch(
|
|
"src.utils.usage_tracking.get_cost_per_token",
|
|
return_value=(input_cost, output_cost, cache_read_cost),
|
|
)
|
|
|
|
usage_tracking.track_usage(
|
|
filename, model_id, input_tokens, output_tokens, cache_creation, cache_read
|
|
)
|
|
|
|
# Verify per-file, per-model data
|
|
usage_data = usage_tracking.get_usage_data()
|
|
file_data = usage_data["per_file_per_model"][filename][model_id]
|
|
|
|
assert file_data["input_tokens"] == input_tokens
|
|
assert file_data["output_tokens"] == output_tokens
|
|
assert file_data["cache_creation_tokens"] == cache_creation
|
|
assert file_data["cache_read_tokens"] == cache_read
|
|
assert file_data["total_tokens"] == input_tokens + output_tokens + cache_creation + cache_read
|
|
assert file_data["request_count"] == 1
|
|
|
|
# Verify cost calculation
|
|
expected_cost = (
|
|
(input_tokens * input_cost)
|
|
+ (cache_creation * input_cost) # Cache creation uses input rate
|
|
+ (cache_read * cache_read_cost)
|
|
+ (output_tokens * output_cost)
|
|
)
|
|
assert abs(file_data["total_cost"] - expected_cost) < 1e-10
|
|
|
|
# Verify global totals
|
|
global_data = usage_data["global"]
|
|
assert global_data["total_fresh_input_tokens"] == input_tokens
|
|
assert global_data["total_cache_tokens"] == cache_creation + cache_read
|
|
assert global_data["total_input_tokens"] == input_tokens + cache_creation + cache_read
|
|
assert global_data["total_output_tokens"] == output_tokens
|
|
assert global_data["total_tokens"] == input_tokens + output_tokens + cache_creation + cache_read
|
|
assert abs(global_data["total_cost"] - expected_cost) < 1e-10
|
|
assert global_data["total_requests"] == 1
|
|
|
|
def test_track_usage_multiple_calls_same_file_model(self, mocker):
|
|
"""Test tracking multiple calls for same file/model (accumulation)."""
|
|
filename = "file1.txt"
|
|
model_id = "anthropic.claude-3-5-sonnet-20240620-v1:0"
|
|
input_cost = 0.003 / 1000.0
|
|
output_cost = 0.015 / 1000.0
|
|
cache_read_cost = 0.0003 / 1000.0
|
|
mocker.patch(
|
|
"src.utils.usage_tracking.get_cost_per_token",
|
|
return_value=(input_cost, output_cost, cache_read_cost),
|
|
)
|
|
|
|
# First call
|
|
usage_tracking.track_usage(filename, model_id, 100, 50, 10, 20)
|
|
# Second call
|
|
usage_tracking.track_usage(filename, model_id, 200, 100, 20, 30)
|
|
|
|
usage_data = usage_tracking.get_usage_data()
|
|
file_data = usage_data["per_file_per_model"][filename][model_id]
|
|
|
|
# Verify accumulation
|
|
assert file_data["input_tokens"] == 300 # 100 + 200
|
|
assert file_data["output_tokens"] == 150 # 50 + 100
|
|
assert file_data["cache_creation_tokens"] == 30 # 10 + 20
|
|
assert file_data["cache_read_tokens"] == 50 # 20 + 30
|
|
assert file_data["total_tokens"] == 530 # 300 + 150 + 30 + 50
|
|
assert file_data["request_count"] == 2
|
|
|
|
# Verify global totals
|
|
global_data = usage_data["global"]
|
|
assert global_data["total_fresh_input_tokens"] == 300
|
|
assert global_data["total_cache_tokens"] == 80 # 30 + 50
|
|
assert global_data["total_input_tokens"] == 380 # 300 + 80
|
|
assert global_data["total_output_tokens"] == 150
|
|
assert global_data["total_tokens"] == 530
|
|
assert global_data["total_requests"] == 2
|
|
|
|
def test_track_usage_multiple_files_models(self, mocker):
|
|
"""Test tracking calls for different files/models (isolation)."""
|
|
input_cost = 0.003 / 1000.0
|
|
output_cost = 0.015 / 1000.0
|
|
cache_read_cost = 0.0003 / 1000.0
|
|
mocker.patch(
|
|
"src.utils.usage_tracking.get_cost_per_token",
|
|
return_value=(input_cost, output_cost, cache_read_cost),
|
|
)
|
|
|
|
# File 1, Model 1
|
|
usage_tracking.track_usage("file1.txt", "model1", 100, 50, 10, 20)
|
|
# File 1, Model 2
|
|
usage_tracking.track_usage("file1.txt", "model2", 200, 100, 20, 30)
|
|
# File 2, Model 1
|
|
usage_tracking.track_usage("file2.txt", "model1", 150, 75, 15, 25)
|
|
|
|
usage_data = usage_tracking.get_usage_data()
|
|
|
|
# Verify file1, model1
|
|
file1_model1 = usage_data["per_file_per_model"]["file1.txt"]["model1"]
|
|
assert file1_model1["input_tokens"] == 100
|
|
assert file1_model1["request_count"] == 1
|
|
|
|
# Verify file1, model2
|
|
file1_model2 = usage_data["per_file_per_model"]["file1.txt"]["model2"]
|
|
assert file1_model2["input_tokens"] == 200
|
|
assert file1_model2["request_count"] == 1
|
|
|
|
# Verify file2, model1
|
|
file2_model1 = usage_data["per_file_per_model"]["file2.txt"]["model1"]
|
|
assert file2_model1["input_tokens"] == 150
|
|
assert file2_model1["request_count"] == 1
|
|
|
|
# Verify global totals (sum of all)
|
|
global_data = usage_data["global"]
|
|
assert global_data["total_fresh_input_tokens"] == 450 # 100 + 200 + 150
|
|
assert global_data["total_cache_tokens"] == 120 # (10+20) + (20+30) + (15+25) = 30 + 50 + 40
|
|
assert global_data["total_requests"] == 3
|
|
|
|
@pytest.mark.parametrize(
|
|
"filename, input_tokens, output_tokens, cache_creation, cache_read, should_reject",
|
|
[
|
|
# Empty filename
|
|
("", 100, 50, 0, 0, True),
|
|
(None, 100, 50, 0, 0, True),
|
|
# Negative input tokens
|
|
("file.txt", -1, 50, 0, 0, True),
|
|
# Negative output tokens
|
|
("file.txt", 100, -1, 0, 0, True),
|
|
# Negative cache creation tokens
|
|
("file.txt", 100, 50, -1, 0, True),
|
|
# Negative cache read tokens
|
|
("file.txt", 100, 50, 0, -1, True),
|
|
# Valid (zero values are OK)
|
|
("file.txt", 0, 0, 0, 0, False),
|
|
],
|
|
)
|
|
def test_track_usage_input_validation(
|
|
self,
|
|
filename,
|
|
input_tokens,
|
|
output_tokens,
|
|
cache_creation,
|
|
cache_read,
|
|
should_reject,
|
|
mocker,
|
|
caplog,
|
|
):
|
|
"""Test input validation (should reject invalid inputs)."""
|
|
mocker.patch(
|
|
"src.utils.usage_tracking.get_cost_per_token",
|
|
return_value=(0.001, 0.002, 0.0001),
|
|
)
|
|
|
|
initial_data = usage_tracking.get_usage_data()
|
|
initial_requests = initial_data["global"]["total_requests"]
|
|
|
|
with caplog.at_level(logging.WARNING):
|
|
usage_tracking.track_usage(
|
|
filename, "model1", input_tokens, output_tokens, cache_creation, cache_read
|
|
)
|
|
|
|
final_data = usage_tracking.get_usage_data()
|
|
final_requests = final_data["global"]["total_requests"]
|
|
|
|
if should_reject:
|
|
# Should not have updated usage data
|
|
assert final_requests == initial_requests
|
|
assert "Invalid usage data" in caplog.text
|
|
else:
|
|
# Should have updated usage data
|
|
assert final_requests == initial_requests + 1
|
|
|
|
def test_track_usage_thread_safety(self, mocker):
|
|
"""Test thread safety with concurrent calls."""
|
|
input_cost = 0.003 / 1000.0
|
|
output_cost = 0.015 / 1000.0
|
|
cache_read_cost = 0.0003 / 1000.0
|
|
mocker.patch(
|
|
"src.utils.usage_tracking.get_cost_per_token",
|
|
return_value=(input_cost, output_cost, cache_read_cost),
|
|
)
|
|
|
|
num_threads = 10
|
|
calls_per_thread = 50
|
|
filename = "file1.txt"
|
|
model_id = "model1"
|
|
|
|
def track_multiple_calls():
|
|
for _ in range(calls_per_thread):
|
|
usage_tracking.track_usage(filename, model_id, 10, 5, 2, 3)
|
|
time.sleep(0.001) # Small delay to increase chance of race condition
|
|
|
|
threads = []
|
|
for _ in range(num_threads):
|
|
thread = threading.Thread(target=track_multiple_calls)
|
|
threads.append(thread)
|
|
thread.start()
|
|
|
|
# Wait for all threads to complete
|
|
for thread in threads:
|
|
thread.join()
|
|
|
|
usage_data = usage_tracking.get_usage_data()
|
|
file_data = usage_data["per_file_per_model"][filename][model_id]
|
|
|
|
# Verify all calls were tracked (no data loss)
|
|
expected_requests = num_threads * calls_per_thread
|
|
assert file_data["request_count"] == expected_requests
|
|
|
|
# Verify token totals are correct
|
|
expected_input = expected_requests * 10
|
|
expected_output = expected_requests * 5
|
|
expected_cache_creation = expected_requests * 2
|
|
expected_cache_read = expected_requests * 3
|
|
|
|
assert file_data["input_tokens"] == expected_input
|
|
assert file_data["output_tokens"] == expected_output
|
|
assert file_data["cache_creation_tokens"] == expected_cache_creation
|
|
assert file_data["cache_read_tokens"] == expected_cache_read
|
|
|
|
# Verify global totals
|
|
global_data = usage_data["global"]
|
|
assert global_data["total_requests"] == expected_requests
|
|
assert global_data["total_fresh_input_tokens"] == expected_input
|
|
assert global_data["total_cache_tokens"] == expected_cache_creation + expected_cache_read
|
|
|
|
def test_track_usage_cost_calculation(self, mocker):
|
|
"""Test that cost calculation is correct for all token types."""
|
|
filename = "file1.txt"
|
|
model_id = "model1"
|
|
input_cost = 0.003 / 1000.0 # $0.003 per 1k = $0.000003 per token
|
|
output_cost = 0.015 / 1000.0 # $0.015 per 1k = $0.000015 per token
|
|
cache_read_cost = 0.0003 / 1000.0 # $0.0003 per 1k = $0.0000003 per token
|
|
mocker.patch(
|
|
"src.utils.usage_tracking.get_cost_per_token",
|
|
return_value=(input_cost, output_cost, cache_read_cost),
|
|
)
|
|
|
|
# Track usage with specific token counts
|
|
input_tokens = 1000
|
|
output_tokens = 500
|
|
cache_creation_tokens = 200 # Charged at input rate
|
|
cache_read_tokens = 300 # Charged at cache read rate
|
|
|
|
usage_tracking.track_usage(
|
|
filename, model_id, input_tokens, output_tokens, cache_creation_tokens, cache_read_tokens
|
|
)
|
|
|
|
usage_data = usage_tracking.get_usage_data()
|
|
file_data = usage_data["per_file_per_model"][filename][model_id]
|
|
|
|
# Calculate expected cost manually
|
|
expected_cost = (
|
|
(input_tokens * input_cost)
|
|
+ (cache_creation_tokens * input_cost) # Same as input rate
|
|
+ (cache_read_tokens * cache_read_cost) # Lower rate
|
|
+ (output_tokens * output_cost)
|
|
)
|
|
|
|
assert abs(file_data["total_cost"] - expected_cost) < 1e-10
|
|
|
|
# Verify cache creation uses input rate, not cache read rate
|
|
cache_creation_cost = cache_creation_tokens * input_cost
|
|
cache_read_cost_calc = cache_read_tokens * cache_read_cost
|
|
assert cache_creation_cost > cache_read_cost_calc # Input rate > cache read rate
|
|
|
|
|
|
class TestGetUsageData:
|
|
"""Tests for get_usage_data() function - Phase 2."""
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def reset_before_test(self):
|
|
"""Reset usage data before each test."""
|
|
usage_tracking.reset_usage_data()
|
|
yield
|
|
usage_tracking.reset_usage_data()
|
|
|
|
def test_get_usage_data_empty_state(self):
|
|
"""Test getting usage data when no usage has been tracked."""
|
|
usage_data = usage_tracking.get_usage_data()
|
|
|
|
assert "per_file_per_model" in usage_data
|
|
assert "global" in usage_data
|
|
assert usage_data["per_file_per_model"] == {}
|
|
assert usage_data["global"]["total_tokens"] == 0
|
|
assert usage_data["global"]["total_cost"] == 0.0
|
|
assert usage_data["global"]["total_requests"] == 0
|
|
|
|
def test_get_usage_data_single_file_single_model(self, mocker):
|
|
"""Test getting usage data with single file and model."""
|
|
input_cost = 0.003 / 1000.0
|
|
output_cost = 0.015 / 1000.0
|
|
cache_read_cost = 0.0003 / 1000.0
|
|
mocker.patch(
|
|
"src.utils.usage_tracking.get_cost_per_token",
|
|
return_value=(input_cost, output_cost, cache_read_cost),
|
|
)
|
|
|
|
usage_tracking.track_usage("file1.txt", "model1", 100, 50, 10, 20)
|
|
usage_data = usage_tracking.get_usage_data()
|
|
|
|
assert "per_file_per_model" in usage_data
|
|
assert "global" in usage_data
|
|
assert "file1.txt" in usage_data["per_file_per_model"]
|
|
assert "model1" in usage_data["per_file_per_model"]["file1.txt"]
|
|
assert usage_data["per_file_per_model"]["file1.txt"]["model1"]["input_tokens"] == 100
|
|
assert usage_data["global"]["total_fresh_input_tokens"] == 100
|
|
|
|
def test_get_usage_data_multiple_files_single_model(self, mocker):
|
|
"""Test getting usage data with multiple files and single model."""
|
|
input_cost = 0.003 / 1000.0
|
|
output_cost = 0.015 / 1000.0
|
|
cache_read_cost = 0.0003 / 1000.0
|
|
mocker.patch(
|
|
"src.utils.usage_tracking.get_cost_per_token",
|
|
return_value=(input_cost, output_cost, cache_read_cost),
|
|
)
|
|
|
|
usage_tracking.track_usage("file1.txt", "model1", 100, 50, 0, 0)
|
|
usage_tracking.track_usage("file2.txt", "model1", 200, 100, 0, 0)
|
|
usage_data = usage_tracking.get_usage_data()
|
|
|
|
assert len(usage_data["per_file_per_model"]) == 2
|
|
assert "file1.txt" in usage_data["per_file_per_model"]
|
|
assert "file2.txt" in usage_data["per_file_per_model"]
|
|
assert usage_data["global"]["total_fresh_input_tokens"] == 300
|
|
|
|
def test_get_usage_data_single_file_multiple_models(self, mocker):
|
|
"""Test getting usage data with single file and multiple models."""
|
|
input_cost = 0.003 / 1000.0
|
|
output_cost = 0.015 / 1000.0
|
|
cache_read_cost = 0.0003 / 1000.0
|
|
mocker.patch(
|
|
"src.utils.usage_tracking.get_cost_per_token",
|
|
return_value=(input_cost, output_cost, cache_read_cost),
|
|
)
|
|
|
|
usage_tracking.track_usage("file1.txt", "model1", 100, 50, 0, 0)
|
|
usage_tracking.track_usage("file1.txt", "model2", 200, 100, 0, 0)
|
|
usage_data = usage_tracking.get_usage_data()
|
|
|
|
assert len(usage_data["per_file_per_model"]) == 1
|
|
assert "file1.txt" in usage_data["per_file_per_model"]
|
|
assert len(usage_data["per_file_per_model"]["file1.txt"]) == 2
|
|
assert "model1" in usage_data["per_file_per_model"]["file1.txt"]
|
|
assert "model2" in usage_data["per_file_per_model"]["file1.txt"]
|
|
|
|
def test_get_usage_data_multiple_files_multiple_models(self, mocker):
|
|
"""Test getting usage data with multiple files and models."""
|
|
input_cost = 0.003 / 1000.0
|
|
output_cost = 0.015 / 1000.0
|
|
cache_read_cost = 0.0003 / 1000.0
|
|
mocker.patch(
|
|
"src.utils.usage_tracking.get_cost_per_token",
|
|
return_value=(input_cost, output_cost, cache_read_cost),
|
|
)
|
|
|
|
usage_tracking.track_usage("file1.txt", "model1", 100, 50, 10, 20)
|
|
usage_tracking.track_usage("file1.txt", "model2", 200, 100, 20, 30)
|
|
usage_tracking.track_usage("file2.txt", "model1", 150, 75, 15, 25)
|
|
usage_data = usage_tracking.get_usage_data()
|
|
|
|
assert len(usage_data["per_file_per_model"]) == 2
|
|
assert usage_data["global"]["total_fresh_input_tokens"] == 450
|
|
assert usage_data["global"]["total_cache_tokens"] == 120
|
|
assert usage_data["global"]["total_requests"] == 3
|
|
|
|
def test_get_usage_data_deep_copy(self, mocker):
|
|
"""Test that modifying returned dict doesn't affect internal state."""
|
|
input_cost = 0.003 / 1000.0
|
|
output_cost = 0.015 / 1000.0
|
|
cache_read_cost = 0.0003 / 1000.0
|
|
mocker.patch(
|
|
"src.utils.usage_tracking.get_cost_per_token",
|
|
return_value=(input_cost, output_cost, cache_read_cost),
|
|
)
|
|
|
|
usage_tracking.track_usage("file1.txt", "model1", 100, 50, 0, 0)
|
|
usage_data = usage_tracking.get_usage_data()
|
|
|
|
# Modify the returned data
|
|
usage_data["per_file_per_model"]["file1.txt"]["model1"]["input_tokens"] = 999
|
|
usage_data["global"]["total_tokens"] = 999
|
|
|
|
# Get fresh data and verify original values are unchanged
|
|
fresh_data = usage_tracking.get_usage_data()
|
|
assert fresh_data["per_file_per_model"]["file1.txt"]["model1"]["input_tokens"] == 100
|
|
assert fresh_data["global"]["total_tokens"] == 150 # 100 + 50
|
|
|
|
def test_get_usage_data_thread_safety(self, mocker):
|
|
"""Test thread safety when getting data while other thread is updating."""
|
|
input_cost = 0.003 / 1000.0
|
|
output_cost = 0.015 / 1000.0
|
|
cache_read_cost = 0.0003 / 1000.0
|
|
mocker.patch(
|
|
"src.utils.usage_tracking.get_cost_per_token",
|
|
return_value=(input_cost, output_cost, cache_read_cost),
|
|
)
|
|
|
|
results = []
|
|
errors = []
|
|
|
|
def track_and_get():
|
|
try:
|
|
for i in range(10):
|
|
usage_tracking.track_usage(f"file{i}.txt", "model1", 10, 5, 0, 0)
|
|
data = usage_tracking.get_usage_data()
|
|
results.append(data["global"]["total_requests"])
|
|
except Exception as e:
|
|
errors.append(e)
|
|
|
|
threads = []
|
|
for _ in range(5):
|
|
thread = threading.Thread(target=track_and_get)
|
|
threads.append(thread)
|
|
thread.start()
|
|
|
|
for thread in threads:
|
|
thread.join()
|
|
|
|
# Verify no errors occurred
|
|
assert len(errors) == 0
|
|
|
|
# Verify final state is correct
|
|
final_data = usage_tracking.get_usage_data()
|
|
assert final_data["global"]["total_requests"] == 50 # 5 threads * 10 calls
|
|
|
|
|
|
class TestResetUsageData:
|
|
"""Tests for reset_usage_data() function - Phase 2."""
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def reset_before_test(self):
|
|
"""Reset usage data before each test."""
|
|
usage_tracking.reset_usage_data()
|
|
yield
|
|
usage_tracking.reset_usage_data()
|
|
|
|
def test_reset_usage_data_after_tracking(self, mocker):
|
|
"""Test reset after tracking some usage."""
|
|
input_cost = 0.003 / 1000.0
|
|
output_cost = 0.015 / 1000.0
|
|
cache_read_cost = 0.0003 / 1000.0
|
|
mocker.patch(
|
|
"src.utils.usage_tracking.get_cost_per_token",
|
|
return_value=(input_cost, output_cost, cache_read_cost),
|
|
)
|
|
|
|
# Track some usage
|
|
usage_tracking.track_usage("file1.txt", "model1", 100, 50, 10, 20)
|
|
usage_tracking.track_usage("file2.txt", "model2", 200, 100, 20, 30)
|
|
|
|
# Verify data exists
|
|
usage_data = usage_tracking.get_usage_data()
|
|
assert usage_data["global"]["total_requests"] == 2
|
|
assert len(usage_data["per_file_per_model"]) == 2
|
|
|
|
# Reset
|
|
usage_tracking.reset_usage_data()
|
|
|
|
# Verify all data is cleared
|
|
usage_data = usage_tracking.get_usage_data()
|
|
assert usage_data["global"]["total_tokens"] == 0
|
|
assert usage_data["global"]["total_cost"] == 0.0
|
|
assert usage_data["global"]["total_requests"] == 0
|
|
assert usage_data["global"]["total_fresh_input_tokens"] == 0
|
|
assert usage_data["global"]["total_cache_tokens"] == 0
|
|
assert usage_data["per_file_per_model"] == {}
|
|
|
|
def test_reset_usage_data_when_empty(self):
|
|
"""Test reset when already empty (should not error)."""
|
|
# Reset when already empty
|
|
usage_tracking.reset_usage_data()
|
|
|
|
# Verify still empty
|
|
usage_data = usage_tracking.get_usage_data()
|
|
assert usage_data["global"]["total_requests"] == 0
|
|
assert usage_data["per_file_per_model"] == {}
|
|
|
|
# Reset again (should not error)
|
|
usage_tracking.reset_usage_data()
|
|
usage_data = usage_tracking.get_usage_data()
|
|
assert usage_data["global"]["total_requests"] == 0
|
|
|
|
def test_reset_usage_data_clears_both_structures(self, mocker):
|
|
"""Test that reset clears both USAGE_DATA and GLOBAL_USAGE."""
|
|
input_cost = 0.003 / 1000.0
|
|
output_cost = 0.015 / 1000.0
|
|
cache_read_cost = 0.0003 / 1000.0
|
|
mocker.patch(
|
|
"src.utils.usage_tracking.get_cost_per_token",
|
|
return_value=(input_cost, output_cost, cache_read_cost),
|
|
)
|
|
|
|
usage_tracking.track_usage("file1.txt", "model1", 100, 50, 10, 20)
|
|
usage_data = usage_tracking.get_usage_data()
|
|
|
|
# Verify both structures have data
|
|
assert len(usage_data["per_file_per_model"]) > 0
|
|
assert usage_data["global"]["total_requests"] > 0
|
|
|
|
# Reset
|
|
usage_tracking.reset_usage_data()
|
|
usage_data = usage_tracking.get_usage_data()
|
|
|
|
# Verify both structures are cleared
|
|
assert usage_data["per_file_per_model"] == {}
|
|
assert usage_data["global"]["total_requests"] == 0
|
|
assert usage_data["global"]["total_tokens"] == 0
|
|
assert usage_data["global"]["total_cost"] == 0.0
|
|
assert usage_data["global"]["total_fresh_input_tokens"] == 0
|
|
assert usage_data["global"]["total_cache_tokens"] == 0
|
|
|
|
def test_reset_usage_data_thread_safety(self, mocker):
|
|
"""Test thread safety of reset operation."""
|
|
input_cost = 0.003 / 1000.0
|
|
output_cost = 0.015 / 1000.0
|
|
cache_read_cost = 0.0003 / 1000.0
|
|
mocker.patch(
|
|
"src.utils.usage_tracking.get_cost_per_token",
|
|
return_value=(input_cost, output_cost, cache_read_cost),
|
|
)
|
|
|
|
errors = []
|
|
|
|
def track_and_reset():
|
|
try:
|
|
for i in range(10):
|
|
usage_tracking.track_usage(f"file{i}.txt", "model1", 10, 5, 0, 0)
|
|
usage_tracking.reset_usage_data()
|
|
except Exception as e:
|
|
errors.append(e)
|
|
|
|
threads = []
|
|
for _ in range(5):
|
|
thread = threading.Thread(target=track_and_reset)
|
|
threads.append(thread)
|
|
thread.start()
|
|
|
|
for thread in threads:
|
|
thread.join()
|
|
|
|
# Verify no errors occurred
|
|
assert len(errors) == 0
|
|
|
|
# Verify final state is reset
|
|
usage_data = usage_tracking.get_usage_data()
|
|
assert usage_data["global"]["total_requests"] == 0
|
|
assert usage_data["per_file_per_model"] == {}
|
|
|
|
|
|
class TestGetUsageDataframes:
|
|
"""Tests for get_usage_dataframes() function - Phase 3."""
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def reset_before_test(self):
|
|
"""Reset usage data before each test."""
|
|
usage_tracking.reset_usage_data()
|
|
yield
|
|
usage_tracking.reset_usage_data()
|
|
|
|
def test_get_usage_dataframes_empty(self):
|
|
"""Test getting DataFrames when no usage data exists."""
|
|
run_id = "test_batch"
|
|
timestamp = "run_20240101_12-00_test_batch"
|
|
|
|
per_file_df, summary_df = usage_tracking.get_usage_dataframes(run_id, timestamp)
|
|
|
|
assert isinstance(per_file_df, pd.DataFrame)
|
|
assert isinstance(summary_df, pd.DataFrame)
|
|
assert per_file_df.empty
|
|
assert not summary_df.empty # Summary should have one row with zeros
|
|
|
|
def test_get_usage_dataframes_single_file_single_model(self, mocker):
|
|
"""Test getting DataFrames with single file and model."""
|
|
input_cost = 0.003 / 1000.0
|
|
output_cost = 0.015 / 1000.0
|
|
cache_read_cost = 0.0003 / 1000.0
|
|
mocker.patch(
|
|
"src.utils.usage_tracking.get_cost_per_token",
|
|
return_value=(input_cost, output_cost, cache_read_cost),
|
|
)
|
|
|
|
usage_tracking.track_usage("file1.txt", "model1", 100, 50, 10, 20)
|
|
|
|
run_id = "test_batch"
|
|
timestamp = "run_20240101_12-00_test_batch"
|
|
per_file_df, summary_df = usage_tracking.get_usage_dataframes(run_id, timestamp)
|
|
|
|
assert isinstance(per_file_df, pd.DataFrame)
|
|
assert isinstance(summary_df, pd.DataFrame)
|
|
assert len(per_file_df) == 1
|
|
assert len(summary_df) == 1
|
|
|
|
# Verify run_id and timestamp are included
|
|
assert per_file_df["run_id"].iloc[0] == run_id
|
|
assert per_file_df["timestamp"].iloc[0] == timestamp
|
|
assert summary_df["run_id"].iloc[0] == run_id
|
|
assert summary_df["timestamp"].iloc[0] == timestamp
|
|
|
|
def test_get_usage_dataframes_multiple_files_models(self, mocker):
|
|
"""Test getting DataFrames with multiple files and models."""
|
|
input_cost = 0.003 / 1000.0
|
|
output_cost = 0.015 / 1000.0
|
|
cache_read_cost = 0.0003 / 1000.0
|
|
mocker.patch(
|
|
"src.utils.usage_tracking.get_cost_per_token",
|
|
return_value=(input_cost, output_cost, cache_read_cost),
|
|
)
|
|
|
|
usage_tracking.track_usage("file1.txt", "model1", 100, 50, 10, 20)
|
|
usage_tracking.track_usage("file1.txt", "model2", 200, 100, 20, 30)
|
|
usage_tracking.track_usage("file2.txt", "model1", 150, 75, 15, 25)
|
|
|
|
run_id = "test_batch"
|
|
timestamp = "run_20240101_12-00_test_batch"
|
|
per_file_df, summary_df = usage_tracking.get_usage_dataframes(run_id, timestamp)
|
|
|
|
assert len(per_file_df) == 3 # 3 file/model combinations
|
|
assert len(summary_df) == 1 # Single summary row
|
|
|
|
def test_get_usage_dataframes_independent(self, mocker):
|
|
"""Test that modifying one DataFrame doesn't affect the other."""
|
|
input_cost = 0.003 / 1000.0
|
|
output_cost = 0.015 / 1000.0
|
|
cache_read_cost = 0.0003 / 1000.0
|
|
mocker.patch(
|
|
"src.utils.usage_tracking.get_cost_per_token",
|
|
return_value=(input_cost, output_cost, cache_read_cost),
|
|
)
|
|
|
|
usage_tracking.track_usage("file1.txt", "model1", 100, 50, 0, 0)
|
|
|
|
run_id = "test_batch"
|
|
timestamp = "run_20240101_12-00_test_batch"
|
|
per_file_df, summary_df = usage_tracking.get_usage_dataframes(run_id, timestamp)
|
|
|
|
# Modify per_file_df
|
|
per_file_df.loc[0, "input_tokens"] = 999
|
|
|
|
# Verify summary_df is unchanged
|
|
assert summary_df["total_fresh_input_tokens"].iloc[0] == 100
|
|
|
|
|
|
class TestGeneratePerFilePerModelDf:
|
|
"""Tests for _generate_per_file_per_model_df() function - Phase 3."""
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def reset_before_test(self):
|
|
"""Reset usage data before each test."""
|
|
usage_tracking.reset_usage_data()
|
|
yield
|
|
usage_tracking.reset_usage_data()
|
|
|
|
def test_generate_per_file_per_model_df_empty(self):
|
|
"""Test generating DataFrame when no usage data exists."""
|
|
run_id = "test_batch"
|
|
timestamp = "run_20240101_12-00_test_batch"
|
|
|
|
df = usage_tracking._generate_per_file_per_model_df(run_id, timestamp)
|
|
|
|
assert isinstance(df, pd.DataFrame)
|
|
assert df.empty
|
|
|
|
def test_generate_per_file_per_model_df_single_file_single_model(self, mocker):
|
|
"""Test generating DataFrame with single file and model."""
|
|
input_cost = 0.003 / 1000.0
|
|
output_cost = 0.015 / 1000.0
|
|
cache_read_cost = 0.0003 / 1000.0
|
|
mocker.patch(
|
|
"src.utils.usage_tracking.get_cost_per_token",
|
|
return_value=(input_cost, output_cost, cache_read_cost),
|
|
)
|
|
|
|
usage_tracking.track_usage("file1.txt", "anthropic.claude-3-5-sonnet-20240620-v1:0", 100, 50, 10, 20)
|
|
|
|
run_id = "test_batch"
|
|
timestamp = "run_20240101_12-00_test_batch"
|
|
df = usage_tracking._generate_per_file_per_model_df(run_id, timestamp)
|
|
|
|
assert len(df) == 1
|
|
assert list(df.columns) == [
|
|
"run_id",
|
|
"timestamp",
|
|
"file_name",
|
|
"model_name",
|
|
"input_tokens",
|
|
"output_tokens",
|
|
"cache_creation_tokens",
|
|
"cache_read_tokens",
|
|
"total_tokens",
|
|
"cost_per_token",
|
|
"total_cost",
|
|
]
|
|
|
|
# Verify values
|
|
assert df["run_id"].iloc[0] == run_id
|
|
assert df["timestamp"].iloc[0] == timestamp
|
|
assert df["file_name"].iloc[0] == "file1.txt"
|
|
assert df["model_name"].iloc[0] == "Claude 3.5 Sonnet" # Human-readable
|
|
assert df["input_tokens"].iloc[0] == 100
|
|
assert df["output_tokens"].iloc[0] == 50
|
|
assert df["cache_creation_tokens"].iloc[0] == 10
|
|
assert df["cache_read_tokens"].iloc[0] == 20
|
|
assert df["total_tokens"].iloc[0] == 180 # 100 + 50 + 10 + 20
|
|
|
|
def test_generate_per_file_per_model_df_multiple_files_single_model(self, mocker):
|
|
"""Test generating DataFrame with multiple files and single model."""
|
|
input_cost = 0.003 / 1000.0
|
|
output_cost = 0.015 / 1000.0
|
|
cache_read_cost = 0.0003 / 1000.0
|
|
mocker.patch(
|
|
"src.utils.usage_tracking.get_cost_per_token",
|
|
return_value=(input_cost, output_cost, cache_read_cost),
|
|
)
|
|
|
|
usage_tracking.track_usage("file1.txt", "model1", 100, 50, 0, 0)
|
|
usage_tracking.track_usage("file2.txt", "model1", 200, 100, 0, 0)
|
|
|
|
run_id = "test_batch"
|
|
timestamp = "run_20240101_12-00_test_batch"
|
|
df = usage_tracking._generate_per_file_per_model_df(run_id, timestamp)
|
|
|
|
assert len(df) == 2
|
|
assert set(df["file_name"].values) == {"file1.txt", "file2.txt"}
|
|
|
|
def test_generate_per_file_per_model_df_single_file_multiple_models(self, mocker):
|
|
"""Test generating DataFrame with single file and multiple models."""
|
|
input_cost = 0.003 / 1000.0
|
|
output_cost = 0.015 / 1000.0
|
|
cache_read_cost = 0.0003 / 1000.0
|
|
mocker.patch(
|
|
"src.utils.usage_tracking.get_cost_per_token",
|
|
return_value=(input_cost, output_cost, cache_read_cost),
|
|
)
|
|
|
|
usage_tracking.track_usage("file1.txt", "model1", 100, 50, 0, 0)
|
|
usage_tracking.track_usage("file1.txt", "model2", 200, 100, 0, 0)
|
|
|
|
run_id = "test_batch"
|
|
timestamp = "run_20240101_12-00_test_batch"
|
|
df = usage_tracking._generate_per_file_per_model_df(run_id, timestamp)
|
|
|
|
assert len(df) == 2
|
|
assert set(df["model_name"].values) == {"Claude Unknown", "Claude Unknown"} # Fallback for unknown models
|
|
|
|
def test_generate_per_file_per_model_df_cost_calculation(self, mocker):
|
|
"""Test cost_per_token calculation and rounding."""
|
|
input_cost = 0.003 / 1000.0
|
|
output_cost = 0.015 / 1000.0
|
|
cache_read_cost = 0.0003 / 1000.0
|
|
mocker.patch(
|
|
"src.utils.usage_tracking.get_cost_per_token",
|
|
return_value=(input_cost, output_cost, cache_read_cost),
|
|
)
|
|
|
|
usage_tracking.track_usage("file1.txt", "model1", 1000, 500, 200, 300)
|
|
|
|
run_id = "test_batch"
|
|
timestamp = "run_20240101_12-00_test_batch"
|
|
df = usage_tracking._generate_per_file_per_model_df(run_id, timestamp)
|
|
|
|
# Verify cost_per_token is calculated correctly (total_cost / total_tokens)
|
|
total_tokens = 1000 + 500 + 200 + 300 # 2000
|
|
total_cost = (
|
|
(1000 * input_cost)
|
|
+ (200 * input_cost) # Cache creation uses input rate
|
|
+ (300 * cache_read_cost)
|
|
+ (500 * output_cost)
|
|
)
|
|
expected_cost_per_token = total_cost / total_tokens
|
|
|
|
# Allow for floating point precision (rounding happens in the function)
|
|
assert abs(df["cost_per_token"].iloc[0] - round(expected_cost_per_token, 8)) < 1e-10
|
|
assert df["cost_per_token"].iloc[0] == round(expected_cost_per_token, 8)
|
|
assert df["total_cost"].iloc[0] == round(total_cost, 6)
|
|
|
|
def test_generate_per_file_per_model_df_skips_zero_requests(self, mocker):
|
|
"""Test that files/models with request_count=0 are skipped."""
|
|
# This is tricky - we need to create a defaultdict entry but with 0 requests
|
|
# In practice, this shouldn't happen, but we test the guard clause
|
|
input_cost = 0.003 / 1000.0
|
|
output_cost = 0.015 / 1000.0
|
|
cache_read_cost = 0.0003 / 1000.0
|
|
mocker.patch(
|
|
"src.utils.usage_tracking.get_cost_per_token",
|
|
return_value=(input_cost, output_cost, cache_read_cost),
|
|
)
|
|
|
|
# Track some usage
|
|
usage_tracking.track_usage("file1.txt", "model1", 100, 50, 0, 0)
|
|
|
|
run_id = "test_batch"
|
|
timestamp = "run_20240101_12-00_test_batch"
|
|
df = usage_tracking._generate_per_file_per_model_df(run_id, timestamp)
|
|
|
|
# Should only have the one entry we tracked
|
|
assert len(df) == 1
|
|
assert df["file_name"].iloc[0] == "file1.txt"
|
|
|
|
|
|
class TestGenerateBatchSummaryDf:
|
|
"""Tests for _generate_batch_summary_df() function - Phase 3."""
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def reset_before_test(self):
|
|
"""Reset usage data before each test."""
|
|
usage_tracking.reset_usage_data()
|
|
yield
|
|
usage_tracking.reset_usage_data()
|
|
|
|
def test_generate_batch_summary_df_empty(self):
|
|
"""Test generating summary DataFrame when no usage data exists."""
|
|
run_id = "test_batch"
|
|
timestamp = "run_20240101_12-00_test_batch"
|
|
|
|
df = usage_tracking._generate_batch_summary_df(run_id, timestamp)
|
|
|
|
assert isinstance(df, pd.DataFrame)
|
|
assert len(df) == 1 # Single row
|
|
assert df["total_tokens"].iloc[0] == 0
|
|
assert df["total_cost"].iloc[0] == 0.0
|
|
|
|
def test_generate_batch_summary_df_single_file(self, mocker):
|
|
"""Test generating summary DataFrame with single file."""
|
|
input_cost = 0.003 / 1000.0
|
|
output_cost = 0.015 / 1000.0
|
|
cache_read_cost = 0.0003 / 1000.0
|
|
mocker.patch(
|
|
"src.utils.usage_tracking.get_cost_per_token",
|
|
return_value=(input_cost, output_cost, cache_read_cost),
|
|
)
|
|
|
|
usage_tracking.track_usage("file1.txt", "model1", 100, 50, 10, 20)
|
|
|
|
run_id = "test_batch"
|
|
timestamp = "run_20240101_12-00_test_batch"
|
|
df = usage_tracking._generate_batch_summary_df(run_id, timestamp)
|
|
|
|
assert len(df) == 1
|
|
assert df["total_fresh_input_tokens"].iloc[0] == 100
|
|
assert df["total_cache_tokens"].iloc[0] == 30 # 10 + 20
|
|
assert df["total_input_tokens"].iloc[0] == 130 # 100 + 30
|
|
assert df["total_output_tokens"].iloc[0] == 50
|
|
assert df["total_tokens"].iloc[0] == 180 # 100 + 50 + 10 + 20
|
|
|
|
# Verify averages (divide by 1 file)
|
|
assert df["average_fresh_input_tokens"].iloc[0] == 100.0
|
|
assert df["average_cache_tokens"].iloc[0] == 30.0
|
|
|
|
def test_generate_batch_summary_df_multiple_files(self, mocker):
|
|
"""Test generating summary DataFrame with multiple files."""
|
|
input_cost = 0.003 / 1000.0
|
|
output_cost = 0.015 / 1000.0
|
|
cache_read_cost = 0.0003 / 1000.0
|
|
mocker.patch(
|
|
"src.utils.usage_tracking.get_cost_per_token",
|
|
return_value=(input_cost, output_cost, cache_read_cost),
|
|
)
|
|
|
|
usage_tracking.track_usage("file1.txt", "model1", 100, 50, 10, 20)
|
|
usage_tracking.track_usage("file2.txt", "model1", 200, 100, 20, 30)
|
|
usage_tracking.track_usage("file3.txt", "model2", 150, 75, 15, 25)
|
|
|
|
run_id = "test_batch"
|
|
timestamp = "run_20240101_12-00_test_batch"
|
|
df = usage_tracking._generate_batch_summary_df(run_id, timestamp)
|
|
|
|
assert len(df) == 1
|
|
# Verify totals
|
|
assert df["total_fresh_input_tokens"].iloc[0] == 450 # 100 + 200 + 150
|
|
assert df["total_cache_tokens"].iloc[0] == 120 # (10+20) + (20+30) + (15+25)
|
|
assert df["total_input_tokens"].iloc[0] == 570 # 450 + 120
|
|
|
|
# Verify averages (divide by 3 files)
|
|
assert df["average_fresh_input_tokens"].iloc[0] == 150.0 # 450 / 3
|
|
assert df["average_cache_tokens"].iloc[0] == 40.0 # 120 / 3
|
|
|
|
def test_generate_batch_summary_df_columns(self, mocker):
|
|
"""Test that summary DataFrame has all required columns."""
|
|
input_cost = 0.003 / 1000.0
|
|
output_cost = 0.015 / 1000.0
|
|
cache_read_cost = 0.0003 / 1000.0
|
|
mocker.patch(
|
|
"src.utils.usage_tracking.get_cost_per_token",
|
|
return_value=(input_cost, output_cost, cache_read_cost),
|
|
)
|
|
|
|
usage_tracking.track_usage("file1.txt", "model1", 100, 50, 0, 0)
|
|
|
|
run_id = "test_batch"
|
|
timestamp = "run_20240101_12-00_test_batch"
|
|
df = usage_tracking._generate_batch_summary_df(run_id, timestamp)
|
|
|
|
expected_columns = [
|
|
"run_id",
|
|
"timestamp",
|
|
"average_fresh_input_tokens",
|
|
"average_cache_tokens",
|
|
"average_input_tokens",
|
|
"average_output_tokens",
|
|
"average_total_tokens",
|
|
"average_cost",
|
|
"total_fresh_input_tokens",
|
|
"total_cache_tokens",
|
|
"total_input_tokens",
|
|
"total_output_tokens",
|
|
"total_tokens",
|
|
"total_cost",
|
|
"region_mode",
|
|
]
|
|
|
|
assert list(df.columns) == expected_columns
|
|
|
|
def test_generate_batch_summary_df_region_mode_single(self, mocker):
|
|
"""Test region_mode is set correctly for single-region mode."""
|
|
input_cost = 0.003 / 1000.0
|
|
output_cost = 0.015 / 1000.0
|
|
cache_read_cost = 0.0003 / 1000.0
|
|
mocker.patch(
|
|
"src.utils.usage_tracking.get_cost_per_token",
|
|
return_value=(input_cost, output_cost, cache_read_cost),
|
|
)
|
|
mocker.patch("src.utils.usage_tracking.config.ENABLE_RUNTIME_ROTATION", False)
|
|
|
|
usage_tracking.track_usage("file1.txt", "model1", 100, 50, 0, 0)
|
|
|
|
run_id = "test_batch"
|
|
timestamp = "run_20240101_12-00_test_batch"
|
|
df = usage_tracking._generate_batch_summary_df(run_id, timestamp)
|
|
|
|
assert df["region_mode"].iloc[0] == "single-region"
|
|
|
|
def test_generate_batch_summary_df_region_mode_cross(self, mocker):
|
|
"""Test region_mode is set correctly for cross-region mode."""
|
|
input_cost = 0.003 / 1000.0
|
|
output_cost = 0.015 / 1000.0
|
|
cache_read_cost = 0.0003 / 1000.0
|
|
mocker.patch(
|
|
"src.utils.usage_tracking.get_cost_per_token",
|
|
return_value=(input_cost, output_cost, cache_read_cost),
|
|
)
|
|
mocker.patch("src.utils.usage_tracking.config.ENABLE_RUNTIME_ROTATION", True)
|
|
|
|
usage_tracking.track_usage("file1.txt", "model1", 100, 50, 0, 0)
|
|
|
|
run_id = "test_batch"
|
|
timestamp = "run_20240101_12-00_test_batch"
|
|
df = usage_tracking._generate_batch_summary_df(run_id, timestamp)
|
|
|
|
assert df["region_mode"].iloc[0] == "cross-region"
|
|
|
|
def test_generate_batch_summary_df_rounding(self, mocker):
|
|
"""Test that averages are properly rounded."""
|
|
input_cost = 0.003 / 1000.0
|
|
output_cost = 0.015 / 1000.0
|
|
cache_read_cost = 0.0003 / 1000.0
|
|
mocker.patch(
|
|
"src.utils.usage_tracking.get_cost_per_token",
|
|
return_value=(input_cost, output_cost, cache_read_cost),
|
|
)
|
|
|
|
# Create scenario where division results in many decimal places
|
|
usage_tracking.track_usage("file1.txt", "model1", 100, 50, 0, 0)
|
|
usage_tracking.track_usage("file2.txt", "model1", 333, 167, 0, 0)
|
|
|
|
run_id = "test_batch"
|
|
timestamp = "run_20240101_12-00_test_batch"
|
|
df = usage_tracking._generate_batch_summary_df(run_id, timestamp)
|
|
|
|
# Verify rounding (to 2 decimal places for averages)
|
|
assert isinstance(df["average_fresh_input_tokens"].iloc[0], float)
|
|
# 433 / 2 = 216.5, should be rounded to 2 decimals
|
|
assert df["average_fresh_input_tokens"].iloc[0] == 216.5
|
|
|
|
def test_generate_batch_summary_df_token_relationship(self, mocker):
|
|
"""Test that total_fresh_input_tokens + total_cache_tokens = total_input_tokens."""
|
|
input_cost = 0.003 / 1000.0
|
|
output_cost = 0.015 / 1000.0
|
|
cache_read_cost = 0.0003 / 1000.0
|
|
mocker.patch(
|
|
"src.utils.usage_tracking.get_cost_per_token",
|
|
return_value=(input_cost, output_cost, cache_read_cost),
|
|
)
|
|
|
|
usage_tracking.track_usage("file1.txt", "model1", 100, 50, 10, 20)
|
|
|
|
run_id = "test_batch"
|
|
timestamp = "run_20240101_12-00_test_batch"
|
|
df = usage_tracking._generate_batch_summary_df(run_id, timestamp)
|
|
|
|
fresh = df["total_fresh_input_tokens"].iloc[0]
|
|
cache = df["total_cache_tokens"].iloc[0]
|
|
total_input = df["total_input_tokens"].iloc[0]
|
|
|
|
assert fresh + cache == total_input
|
|
|
|
|
|
class TestExportUsageCostCsvsLocal:
|
|
"""Tests for export_usage_cost_csvs_local() function - Phase 3."""
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def reset_before_test(self):
|
|
"""Reset usage data before each test."""
|
|
usage_tracking.reset_usage_data()
|
|
yield
|
|
usage_tracking.reset_usage_data()
|
|
|
|
def test_export_usage_cost_csvs_local_with_data(self, mocker, tmp_path, caplog):
|
|
"""Test exporting CSV files with valid usage data."""
|
|
input_cost = 0.003 / 1000.0
|
|
output_cost = 0.015 / 1000.0
|
|
cache_read_cost = 0.0003 / 1000.0
|
|
mocker.patch(
|
|
"src.utils.usage_tracking.get_cost_per_token",
|
|
return_value=(input_cost, output_cost, cache_read_cost),
|
|
)
|
|
mocker.patch(
|
|
"src.utils.usage_tracking.config.CONSOLIDATED_OUTPUT_DIRECTORY", str(tmp_path)
|
|
)
|
|
|
|
usage_tracking.track_usage("file1.txt", "model1", 100, 50, 10, 20)
|
|
|
|
run_id = "test_batch"
|
|
run_timestamp = "run_20240101_12-00_test_batch"
|
|
|
|
with caplog.at_level(logging.INFO):
|
|
usage_tracking.export_usage_cost_csvs_local(run_id, run_timestamp)
|
|
|
|
# Verify directory was created
|
|
output_dir = tmp_path / run_timestamp
|
|
assert output_dir.exists()
|
|
|
|
# Verify both files were created
|
|
usage_file = output_dir / f"{run_id}-USAGE.csv"
|
|
summary_file = output_dir / f"{run_id}-USAGE-SUMMARY.csv"
|
|
|
|
assert usage_file.exists()
|
|
assert summary_file.exists()
|
|
|
|
# Verify file contents can be read back
|
|
usage_df = pd.read_csv(usage_file)
|
|
summary_df = pd.read_csv(summary_file)
|
|
|
|
assert len(usage_df) == 1
|
|
assert len(summary_df) == 1
|
|
assert usage_df["file_name"].iloc[0] == "file1.txt"
|
|
|
|
# Verify logging
|
|
assert "Saved usage data to:" in caplog.text
|
|
assert "Saved usage summary to:" in caplog.text
|
|
|
|
def test_export_usage_cost_csvs_local_empty_data(self, mocker, tmp_path, caplog):
|
|
"""Test exporting CSV files with empty usage data."""
|
|
mocker.patch(
|
|
"src.utils.usage_tracking.config.CONSOLIDATED_OUTPUT_DIRECTORY", str(tmp_path)
|
|
)
|
|
|
|
run_id = "test_batch"
|
|
run_timestamp = "run_20240101_12-00_test_batch"
|
|
|
|
with caplog.at_level(logging.WARNING):
|
|
usage_tracking.export_usage_cost_csvs_local(run_id, run_timestamp)
|
|
|
|
# Verify directory was created
|
|
output_dir = tmp_path / run_timestamp
|
|
assert output_dir.exists()
|
|
|
|
# Verify files were NOT created (empty data)
|
|
usage_file = output_dir / f"{run_id}-USAGE.csv"
|
|
summary_file = output_dir / f"{run_id}-USAGE-SUMMARY.csv"
|
|
|
|
assert not usage_file.exists()
|
|
assert summary_file.exists() # Summary should still exist (with zeros)
|
|
|
|
# Verify logging
|
|
assert "No usage data to export" in caplog.text
|
|
|
|
def test_export_usage_cost_csvs_local_directory_creation(self, mocker, tmp_path):
|
|
"""Test that directory is created if it doesn't exist."""
|
|
input_cost = 0.003 / 1000.0
|
|
output_cost = 0.015 / 1000.0
|
|
cache_read_cost = 0.0003 / 1000.0
|
|
mocker.patch(
|
|
"src.utils.usage_tracking.get_cost_per_token",
|
|
return_value=(input_cost, output_cost, cache_read_cost),
|
|
)
|
|
mocker.patch(
|
|
"src.utils.usage_tracking.config.CONSOLIDATED_OUTPUT_DIRECTORY", str(tmp_path)
|
|
)
|
|
|
|
usage_tracking.track_usage("file1.txt", "model1", 100, 50, 0, 0)
|
|
|
|
run_id = "test_batch"
|
|
run_timestamp = "run_20240101_12-00_test_batch"
|
|
|
|
# Directory shouldn't exist yet
|
|
output_dir = tmp_path / run_timestamp
|
|
assert not output_dir.exists()
|
|
|
|
usage_tracking.export_usage_cost_csvs_local(run_id, run_timestamp)
|
|
|
|
# Directory should now exist
|
|
assert output_dir.exists()
|
|
|
|
def test_export_usage_cost_csvs_local_file_paths(self, mocker, tmp_path):
|
|
"""Test that file paths are correct."""
|
|
input_cost = 0.003 / 1000.0
|
|
output_cost = 0.015 / 1000.0
|
|
cache_read_cost = 0.0003 / 1000.0
|
|
mocker.patch(
|
|
"src.utils.usage_tracking.get_cost_per_token",
|
|
return_value=(input_cost, output_cost, cache_read_cost),
|
|
)
|
|
mocker.patch(
|
|
"src.utils.usage_tracking.config.CONSOLIDATED_OUTPUT_DIRECTORY", str(tmp_path)
|
|
)
|
|
|
|
usage_tracking.track_usage("file1.txt", "model1", 100, 50, 0, 0)
|
|
|
|
run_id = "test_batch"
|
|
run_timestamp = "run_20240101_12-00_test_batch"
|
|
|
|
usage_tracking.export_usage_cost_csvs_local(run_id, run_timestamp)
|
|
|
|
# Verify exact file paths
|
|
expected_usage_path = tmp_path / run_timestamp / f"{run_id}-USAGE.csv"
|
|
expected_summary_path = tmp_path / run_timestamp / f"{run_id}-USAGE-SUMMARY.csv"
|
|
|
|
assert expected_usage_path.exists()
|
|
assert expected_summary_path.exists()
|
|
|
|
|
|
class TestGetModelName:
|
|
"""Tests for get_model_name() function - Phase 4."""
|
|
|
|
@pytest.mark.parametrize(
|
|
"model_id, expected_name",
|
|
[
|
|
# Claude 2
|
|
("anthropic.claude-instant-v1", "Claude 2"),
|
|
("anthropic.claude-2", "Claude 2"), # Note: claude-v2 doesn't match, needs claude-2
|
|
("ANTHROPIC.CLAUDE-INSTANT-V1", "Claude 2"), # Uppercase
|
|
("Anthropic.Claude-2", "Claude 2"), # Mixed case
|
|
# Claude 3 Haiku
|
|
("anthropic.claude-3-haiku-20240307-v1:0", "Claude 3 Haiku"),
|
|
("anthropic.claude-3.5-haiku-20240307-v1:0", "Claude 3 Haiku"),
|
|
("ANTHROPIC.CLAUDE-3-HAIKU-20240307-V1:0", "Claude 3 Haiku"), # Uppercase
|
|
# Claude 3.5 Sonnet
|
|
("anthropic.claude-3-5-sonnet-20240620-v1:0", "Claude 3.5 Sonnet"),
|
|
("anthropic.claude-3-5-sonnet-20241022-v2:0", "Claude 3.5 Sonnet"),
|
|
("anthropic.claude-3.5-sonnet-20240620-v1:0", "Claude 3.5 Sonnet"),
|
|
("Anthropic.Claude-3-5-Sonnet-20240620-V1:0", "Claude 3.5 Sonnet"), # Mixed case
|
|
# Claude 3.7 Sonnet
|
|
("us.anthropic.claude-3-7-sonnet-20250219-v1:0", "Claude 3.7 Sonnet"),
|
|
("us.anthropic.claude-3.7-sonnet-20250219-v1:0", "Claude 3.7 Sonnet"),
|
|
("US.ANTHROPIC.CLAUDE-3-7-SONNET-20250219-V1:0", "Claude 3.7 Sonnet"), # Uppercase
|
|
# Claude 4 Sonnet
|
|
("us.anthropic.claude-sonnet-4-20250514-v1:0", "Claude 4 Sonnet"),
|
|
("us.anthropic.claude-4-sonnet-20250514-v1:0", "Claude 4 Sonnet"),
|
|
("US.ANTHROPIC.CLAUDE-SONNET-4-20250514-V1:0", "Claude 4 Sonnet"), # Uppercase
|
|
# Claude 4.5 Sonnet
|
|
("us.anthropic.claude-sonnet-4-5-20250929-v1:0", "Claude 4.5 Sonnet"),
|
|
("us.anthropic.claude-sonnet-4.5-20250929-v1:0", "Claude 4.5 Sonnet"),
|
|
# Note: claude-4-5-sonnet and claude-4.5-sonnet don't match the pattern
|
|
# They need claude-sonnet-4 or claude-4-sonnet structure
|
|
("US.ANTHROPIC.CLAUDE-SONNET-4-5-20250929-V1:0", "Claude 4.5 Sonnet"), # Uppercase
|
|
],
|
|
)
|
|
def test_get_model_name_valid_models(self, model_id, expected_name):
|
|
"""Test model name conversion for all supported models."""
|
|
result = usage_tracking.get_model_name(model_id)
|
|
assert result == expected_name
|
|
|
|
@pytest.mark.parametrize(
|
|
"model_id, expected_pattern",
|
|
[
|
|
# Unknown models should use fallback
|
|
("unknown.model.id", "Claude "), # Should extract something or "Unknown"
|
|
("anthropic.claude-999", "Claude "), # Unknown version
|
|
("some.random.model", "Claude "), # Completely unknown
|
|
("", "Claude Unknown"), # Empty string
|
|
],
|
|
)
|
|
def test_get_model_name_unknown_models(self, model_id, expected_pattern):
|
|
"""Test fallback behavior for unknown model IDs."""
|
|
result = usage_tracking.get_model_name(model_id)
|
|
assert result.startswith("Claude ")
|
|
# Should either extract a version or say "Unknown"
|
|
assert "Unknown" in result or result != "Claude "
|
|
|
|
def test_get_model_name_case_insensitivity(self):
|
|
"""Test that function handles case variations correctly."""
|
|
# Test various case combinations
|
|
test_cases = [
|
|
("anthropic.claude-3-5-sonnet-20240620-v1:0", "Claude 3.5 Sonnet"),
|
|
("ANTHROPIC.CLAUDE-3-5-SONNET-20240620-V1:0", "Claude 3.5 Sonnet"),
|
|
("Anthropic.Claude-3-5-Sonnet-20240620-V1:0", "Claude 3.5 Sonnet"),
|
|
("AnThRoPiC.cLaUdE-3-5-sOnNeT-20240620-v1:0", "Claude 3.5 Sonnet"),
|
|
]
|
|
|
|
for model_id, expected in test_cases:
|
|
result = usage_tracking.get_model_name(model_id)
|
|
assert result == expected
|
|
|
|
def test_get_model_name_malformed_ids(self):
|
|
"""Test handling of malformed model IDs."""
|
|
# Test various malformed inputs
|
|
malformed_ids = [
|
|
"claude", # Too short
|
|
".", # Just a dot
|
|
"anthropic.", # Ends with dot
|
|
".claude", # Starts with dot
|
|
"anthropic.claude-3-5", # Missing version info
|
|
]
|
|
|
|
for model_id in malformed_ids:
|
|
result = usage_tracking.get_model_name(model_id)
|
|
# Should not crash and should return something starting with "Claude"
|
|
assert isinstance(result, str)
|
|
assert result.startswith("Claude")
|
|
|
|
def test_get_model_name_special_characters(self):
|
|
"""Test handling of special characters in model IDs."""
|
|
# Test with special characters that might appear
|
|
special_cases = [
|
|
("anthropic.claude-3-5-sonnet-20240620-v1:0", "Claude 3.5 Sonnet"),
|
|
("us.anthropic.claude-sonnet-4-5-20250929-v1:0", "Claude 4.5 Sonnet"),
|
|
]
|
|
|
|
for model_id, expected in special_cases:
|
|
result = usage_tracking.get_model_name(model_id)
|
|
assert result == expected
|
|
|
|
def test_get_model_name_fallback_extraction(self):
|
|
"""Test fallback extraction logic for unknown models."""
|
|
# Test models that should trigger fallback extraction
|
|
# Fallback extracts: model_id.split('.')[-1].split('-')[0]
|
|
test_cases = [
|
|
("anthropic.claude-5-sonnet-v1:0", "Claude claude"), # Extracts "claude" from "claude-5-sonnet-v1:0"
|
|
("anthropic.claude-10-v1:0", "Claude claude"), # Extracts "claude" from "claude-10-v1:0"
|
|
("some.prefix.model-7-v1:0", "Claude model"), # Extracts "model" from "model-7-v1:0"
|
|
]
|
|
|
|
for model_id, expected in test_cases:
|
|
result = usage_tracking.get_model_name(model_id)
|
|
assert result == expected
|
|
|
|
def test_get_model_name_no_dot_fallback(self):
|
|
"""Test fallback when model ID has no dot."""
|
|
# Model IDs without dots should return "Claude Unknown"
|
|
# BUT: if they match known patterns (like "claude-3-5-sonnet"), they still work
|
|
test_cases = [
|
|
("claude-3-5-sonnet", "Claude 3.5 Sonnet"), # Matches pattern even without dot
|
|
("claude-instant-v1", "Claude 2"), # Matches pattern
|
|
("anthropic-claude-3", "Claude Unknown"), # No dot, doesn't match pattern
|
|
("just-text", "Claude Unknown"), # No dot, doesn't match pattern
|
|
]
|
|
|
|
for model_id, expected in test_cases:
|
|
result = usage_tracking.get_model_name(model_id)
|
|
assert result == expected
|
|
|
|
def test_get_model_name_all_supported_models(self):
|
|
"""Test all officially supported model IDs."""
|
|
supported_models = {
|
|
"anthropic.claude-instant-v1": "Claude 2",
|
|
"anthropic.claude-2": "Claude 2", # Note: claude-v2 doesn't match pattern
|
|
"anthropic.claude-3-haiku-20240307-v1:0": "Claude 3 Haiku",
|
|
"anthropic.claude-3-5-sonnet-20240620-v1:0": "Claude 3.5 Sonnet",
|
|
"anthropic.claude-3-5-sonnet-20241022-v2:0": "Claude 3.5 Sonnet",
|
|
"us.anthropic.claude-3-7-sonnet-20250219-v1:0": "Claude 3.7 Sonnet",
|
|
"us.anthropic.claude-sonnet-4-20250514-v1:0": "Claude 4 Sonnet",
|
|
"us.anthropic.claude-sonnet-4-5-20250929-v1:0": "Claude 4.5 Sonnet",
|
|
}
|
|
|
|
for model_id, expected_name in supported_models.items():
|
|
result = usage_tracking.get_model_name(model_id)
|
|
assert result == expected_name, f"Failed for {model_id}: got {result}, expected {expected_name}"
|
|
|