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

1"""Tests for the CostBudget session tracker.""" 

2 

3from __future__ import annotations 

4 

5import threading 

6 

7import pytest 

8from assertpy import assert_that 

9 

10from lintro.ai.budget import CostBudget 

11from lintro.ai.exceptions import AIError 

12 

13# -- Defaults ---------------------------------------------------------------- 

14 

15 

16def test_default_budget_has_no_limit() -> None: 

17 """A budget with no max_cost_usd has unlimited remaining.""" 

18 budget = CostBudget() 

19 assert_that(budget.max_cost_usd).is_none() 

20 assert_that(budget.spent).is_equal_to(0.0) 

21 assert_that(budget.remaining).is_none() 

22 

23 

24def test_budget_with_limit() -> None: 

25 """A budget with max_cost_usd reports remaining correctly.""" 

26 budget = CostBudget(max_cost_usd=5.0) 

27 assert_that(budget.max_cost_usd).is_equal_to(5.0) 

28 assert_that(budget.remaining).is_equal_to(5.0) 

29 

30 

31# -- Record ------------------------------------------------------------------ 

32 

33 

34def test_record_increments_spent() -> None: 

35 """Recording cost increments the spent total.""" 

36 budget = CostBudget(max_cost_usd=10.0) 

37 budget.record(1.5) 

38 assert_that(budget.spent).is_equal_to(1.5) 

39 budget.record(2.0) 

40 assert_that(budget.spent).is_equal_to(3.5) 

41 

42 

43def test_record_updates_remaining() -> None: 

44 """Recording cost decreases remaining budget.""" 

45 budget = CostBudget(max_cost_usd=5.0) 

46 budget.record(3.0) 

47 assert_that(budget.remaining).is_equal_to(2.0) 

48 

49 

50def test_remaining_never_negative() -> None: 

51 """Remaining is clamped to 0.0 when overspent.""" 

52 budget = CostBudget(max_cost_usd=1.0) 

53 budget.record(2.0) 

54 assert_that(budget.remaining).is_equal_to(0.0) 

55 

56 

57# -- Check ------------------------------------------------------------------- 

58 

59 

60def test_check_passes_when_under_limit() -> None: 

61 """check() does not raise when spent is below the limit.""" 

62 budget = CostBudget(max_cost_usd=5.0) 

63 budget.record(2.0) 

64 budget.check() # should not raise 

65 

66 

67def test_check_passes_with_no_limit() -> None: 

68 """check() never raises when max_cost_usd is None.""" 

69 budget = CostBudget() 

70 budget.record(1000.0) 

71 budget.check() # should not raise 

72 

73 

74def test_check_raises_when_at_limit() -> None: 

75 """check() raises AIError when spent equals the limit.""" 

76 budget = CostBudget(max_cost_usd=2.0) 

77 budget.record(2.0) 

78 with pytest.raises(AIError, match="cost budget exceeded"): 

79 budget.check() 

80 

81 

82def test_check_raises_when_over_limit() -> None: 

83 """check() raises AIError when spent exceeds the limit.""" 

84 budget = CostBudget(max_cost_usd=1.0) 

85 budget.record(1.5) 

86 with pytest.raises(AIError, match="cost budget exceeded"): 

87 budget.check() 

88 

89 

90def test_check_error_message_contains_amounts() -> None: 

91 """The AIError message includes both spent and limit amounts.""" 

92 budget = CostBudget(max_cost_usd=2.0) 

93 budget.record(2.5) 

94 with pytest.raises(AIError, match=r"\$2\.5000 spent.*\$2\.00"): 

95 budget.check() 

96 

97 

98# -- Thread safety ----------------------------------------------------------- 

99 

100 

101def test_thread_safety_concurrent_records() -> None: 

102 """Concurrent record() calls produce correct total.""" 

103 budget = CostBudget(max_cost_usd=None) 

104 num_threads = 10 

105 increments_per_thread = 100 

106 cost_per_increment = 0.01 

107 

108 def worker() -> None: 

109 for _ in range(increments_per_thread): 

110 budget.record(cost_per_increment) 

111 

112 threads = [threading.Thread(target=worker) for _ in range(num_threads)] 

113 for t in threads: 

114 t.start() 

115 for t in threads: 

116 t.join() 

117 

118 expected = num_threads * increments_per_thread * cost_per_increment 

119 assert_that(budget.spent).is_close_to(expected, tolerance=1e-9)