"""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}"