Coverage for tests / unit / ai / test_cost.py: 100%

1"""Tests for AI cost estimation.""" 

2 

3from __future__ import annotations 

4 

5from unittest.mock import patch 

6 

7import pytest 

8from assertpy import assert_that 

9 

10from lintro.ai.cost import ( 

11 estimate_cost, 

12 format_cost, 

13 format_token_count, 

14) 

15from lintro.ai.registry import DEFAULT_PRICING, PROVIDERS 

16 

17 

18def test_cost_known_model(): 

19 """Verify cost estimation uses correct pricing for a known model.""" 

20 pricing = PROVIDERS.model_pricing["gpt-4o"] 

21 cost = estimate_cost("gpt-4o", 1000, 500) 

22 expected = (1000 / 1_000_000) * pricing.input_per_million + ( 

23 500 / 1_000_000 

24 ) * pricing.output_per_million 

25 assert_that(cost).is_close_to(expected, 1e-10) 

26 

27 

28def test_cost_unknown_model_uses_default(): 

29 """Verify cost estimation falls back to default pricing for unknown models.""" 

30 cost = estimate_cost("unknown-model", 1000, 500) 

31 expected = (1000 / 1_000_000) * DEFAULT_PRICING.input_per_million + ( 

32 500 / 1_000_000 

33 ) * DEFAULT_PRICING.output_per_million 

34 assert_that(cost).is_close_to(expected, 1e-10) 

35 

36 

37@patch("lintro.ai.cost.logger") 

38def test_cost_unknown_model_logs_debug(mock_logger): 

39 """Verify a debug message is logged for unknown model default pricing.""" 

40 estimate_cost("totally-unknown-model", 100, 50) 

41 mock_logger.debug.assert_called_once() 

42 call_args = mock_logger.debug.call_args[0][0] 

43 assert_that(call_args).contains("totally-unknown-model") 

44 assert_that(call_args).contains("default pricing") 

45 

46 

47@patch("lintro.ai.cost.logger") 

48def test_cost_known_model_does_not_log(mock_logger): 

49 """Verify no debug message is logged when pricing is found for a known model.""" 

50 estimate_cost("gpt-4o", 100, 50) 

51 mock_logger.debug.assert_not_called() 

52 

53 

54def test_cost_zero_tokens(): 

55 """Verify cost is zero when both input and output token counts are zero.""" 

56 cost = estimate_cost("gpt-4o", 0, 0) 

57 assert_that(cost).is_equal_to(0.0) 

58 

59 

60def test_cost_large_token_count(): 

61 """Verify cost estimation handles large token counts correctly.""" 

62 cost = estimate_cost("gpt-4o", 1_000_000, 1_000_000) 

63 assert_that(cost).is_greater_than(0) 

64 

65 

66@pytest.mark.parametrize("model", list(PROVIDERS.model_pricing.keys())) 

67def test_cost_all_known_models_have_pricing(model): 

68 """Verify every registered model produces a positive cost estimate.""" 

69 cost = estimate_cost(model, 1000, 1000) 

70 assert_that(cost).is_greater_than(0) 

71 

72 

73def test_cost_format_small(): 

74 """Verify very small costs are formatted as less-than-threshold.""" 

75 result = format_cost(0.0001) 

76 assert_that(result).is_equal_to("<$0.001") 

77 

78 

79def test_cost_format_normal(): 

80 """Verify normal cost values are formatted with dollar sign and three decimals.""" 

81 result = format_cost(0.005) 

82 assert_that(result).is_equal_to("$0.005") 

83 

84 

85def test_cost_format_larger(): 

86 """Verify larger cost values are formatted correctly with dollar sign.""" 

87 result = format_cost(1.234) 

88 assert_that(result).is_equal_to("$1.234") 

89 

90 

91def test_cost_format_token_count_small(): 

92 """Verify small token counts are formatted with a tilde prefix.""" 

93 result = format_token_count(100) 

94 assert_that(result).is_equal_to("~100") 

95 

96 

97def test_cost_format_token_count_large(): 

98 """Verify large token counts are formatted with comma separators.""" 

99 result = format_token_count(1234567) 

100 assert_that(result).is_equal_to("~1,234,567")