31 lines
1.0 KiB
Python
31 lines
1.0 KiB
Python
from src import cost_tracker
|
|
|
|
def test_estimate_cost():
|
|
# Test unknown model
|
|
assert cost_tracker.estimate_cost("unknown-model", 1000, 1000) == 0.0
|
|
|
|
# Test empty model
|
|
assert cost_tracker.estimate_cost("", 1000, 1000) == 0.0
|
|
|
|
# Test Gemini 3.1 Pro Preview
|
|
# input: 3.50 per M, output: 10.50 per M
|
|
# 1M input + 1M output = 14.00
|
|
cost = cost_tracker.estimate_cost("gemini-3.1-pro-preview", 1_000_000, 1_000_000)
|
|
assert abs(cost - 14.00) < 0.0001
|
|
|
|
# Test Claude Sonnet
|
|
# input: 3.0 per M, output: 15.0 per M
|
|
# 100k input + 10k output = 0.3 + 0.15 = 0.45
|
|
cost = cost_tracker.estimate_cost("claude-3-5-sonnet-20241022", 100_000, 10_000)
|
|
assert abs(cost - 0.45) < 0.0001
|
|
|
|
# Test DeepSeek V3
|
|
# input: 0.27 per M, output: 1.10 per M
|
|
# 1M input + 1M output = 1.37
|
|
cost = cost_tracker.estimate_cost("deepseek-v3", 1_000_000, 1_000_000)
|
|
assert abs(cost - 1.37) < 0.0001
|
|
|
|
if __name__ == "__main__":
|
|
test_estimate_cost()
|
|
print("All cost_tracker tests passed!")
|