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

1"""Tests for token budget estimation and truncation utilities.""" 

2 

3from __future__ import annotations 

4 

5import pytest 

6from assertpy import assert_that 

7 

8from lintro.ai.token_budget import estimate_tokens, truncate_to_budget 

9 

10# -- estimate_tokens --------------------------------------------------------- 

11 

12 

13def test_empty_string_returns_zero() -> None: 

14 """An empty string produces 0 tokens.""" 

15 assert_that(estimate_tokens("")).is_equal_to(0) 

16 

17 

18def test_short_string_returns_at_least_one() -> None: 

19 """A non-empty string shorter than 4 chars still returns 1.""" 

20 assert_that(estimate_tokens("hi")).is_equal_to(1) 

21 

22 

23def test_four_chars_equals_one_token() -> None: 

24 """4 characters maps to exactly 1 token.""" 

25 assert_that(estimate_tokens("abcd")).is_equal_to(1) 

26 

27 

28def test_eight_chars_equals_two_tokens() -> None: 

29 """8 characters maps to 2 tokens.""" 

30 assert_that(estimate_tokens("abcdefgh")).is_equal_to(2) 

31 

32 

33# -- truncate_to_budget ------------------------------------------------------ 

34 

35 

36def test_no_truncation_when_under_budget() -> None: 

37 """Text under the budget is returned as-is with truncated=False.""" 

38 text = "hello" 

39 result, truncated = truncate_to_budget(text, max_tokens=100) 

40 assert_that(result).is_equal_to(text) 

41 assert_that(truncated).is_false() 

42 

43 

44def test_truncation_sets_flag() -> None: 

45 """Text over budget returns truncated=True.""" 

46 text = "a" * 100 

47 result, truncated = truncate_to_budget(text, max_tokens=5) 

48 assert_that(truncated).is_true() 

49 assert_that(len(result)).is_less_than(len(text)) 

50 

51 

52def test_truncation_cuts_at_newline() -> None: 

53 """Truncation prefers cutting at a newline boundary.""" 

54 text = "line1\nline2\nline3\nline4\nline5\n" 

55 result, truncated = truncate_to_budget(text, max_tokens=3) 

56 # max_tokens=3 -> max_chars=12 -> cuts at newline boundary 

57 assert_that(truncated).is_true() 

58 assert_that(result).contains("\n") 

59 # Result should not contain later lines 

60 assert_that(result).does_not_contain("line3") 

61 assert_that(result).does_not_contain("line4") 

62 

63 

64def test_empty_string_no_truncation() -> None: 

65 """An empty string is never truncated.""" 

66 result, truncated = truncate_to_budget("", max_tokens=10) 

67 assert_that(result).is_equal_to("") 

68 assert_that(truncated).is_false() 

69 

70 

71@pytest.mark.parametrize("max_tokens", [0, -1, -5]) 

72def test_truncate_raises_on_non_positive_max_tokens(max_tokens: int) -> None: 

73 """truncate_to_budget raises ValueError for non-positive max_tokens.""" 

74 with pytest.raises(ValueError, match="max_tokens must be positive"): 

75 truncate_to_budget("some text", max_tokens=max_tokens)