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!")