116 lines
3.6 KiB
Python
116 lines
3.6 KiB
Python
"""Tests for utility functions."""
|
|
|
|
from ra_aid.text.processing import truncate_output
|
|
|
|
|
|
def test_normal_truncation():
|
|
"""Test normal truncation behavior with more lines than max."""
|
|
# Create input with 10 lines
|
|
input_lines = [f"Line {i}\n" for i in range(10)]
|
|
input_text = "".join(input_lines)
|
|
|
|
# Truncate to 5 lines
|
|
result = truncate_output(input_text, max_lines=5)
|
|
|
|
# Verify truncation message and content
|
|
assert "[5 lines of output truncated]" in result
|
|
assert "Line 5\n" in result
|
|
assert "Line 9\n" in result
|
|
assert "Line 0\n" not in result
|
|
assert "Line 4\n" not in result
|
|
|
|
|
|
def test_no_truncation_needed():
|
|
"""Test when input is shorter than max_lines."""
|
|
input_text = "Line 1\nLine 2\nLine 3\n"
|
|
result = truncate_output(input_text, max_lines=5)
|
|
|
|
# Should return original text unchanged
|
|
assert result == input_text
|
|
assert "[lines of output truncated]" not in result
|
|
|
|
|
|
def test_empty_input():
|
|
"""Test with empty input."""
|
|
assert truncate_output("") == ""
|
|
assert truncate_output(None) == ""
|
|
|
|
|
|
def test_exact_max_lines():
|
|
"""Test when input is exactly max_lines."""
|
|
# Create input with exactly 5 lines
|
|
input_lines = [f"Line {i}\n" for i in range(5)]
|
|
input_text = "".join(input_lines)
|
|
|
|
result = truncate_output(input_text, max_lines=5)
|
|
|
|
# Should return original text unchanged
|
|
assert result == input_text
|
|
assert "[lines of output truncated]" not in result
|
|
|
|
|
|
def test_different_line_endings():
|
|
"""Test with different line endings (\\n, \\r\\n, \\r)."""
|
|
# Mix of different line endings
|
|
input_text = "Line 1\nLine 2\r\nLine 3\rLine 4\nLine 5\r\nLine 6"
|
|
|
|
result = truncate_output(input_text, max_lines=3)
|
|
|
|
# Should preserve line endings in truncated output
|
|
assert "[3 lines of output truncated]" in result
|
|
assert "Line 4" in result
|
|
assert "Line 6" in result
|
|
assert "Line 1" not in result
|
|
|
|
|
|
def test_ansi_sequences():
|
|
"""Test with ANSI escape sequences."""
|
|
# Input with ANSI color codes
|
|
input_lines = [
|
|
"\033[31mRed Line 1\033[0m\n",
|
|
"\033[32mGreen Line 2\033[0m\n",
|
|
"\033[34mBlue Line 3\033[0m\n",
|
|
"\033[33mYellow Line 4\033[0m\n",
|
|
]
|
|
input_text = "".join(input_lines)
|
|
|
|
result = truncate_output(input_text, max_lines=2)
|
|
|
|
# Should preserve ANSI sequences in truncated output
|
|
assert "[2 lines of output truncated]" in result
|
|
assert "\033[34mBlue Line 3\033[0m" in result
|
|
assert "\033[33mYellow Line 4\033[0m" in result
|
|
assert "\033[31mRed Line 1\033[0m" not in result
|
|
|
|
|
|
def test_custom_max_lines():
|
|
"""Test with custom max_lines value."""
|
|
# Create input with 100 lines
|
|
input_lines = [f"Line {i}\n" for i in range(100)]
|
|
input_text = "".join(input_lines)
|
|
|
|
# Test with custom max_lines=10
|
|
result = truncate_output(input_text, max_lines=10)
|
|
|
|
# Should have truncation message and last 10 lines
|
|
assert "[90 lines of output truncated]" in result
|
|
assert "Line 90\n" in result
|
|
assert "Line 99\n" in result
|
|
assert "Line 0\n" not in result
|
|
assert "Line 89\n" not in result
|
|
|
|
|
|
def test_no_trailing_newline():
|
|
"""Test with input that doesn't end in newline."""
|
|
input_lines = [f"Line {i}" for i in range(10)]
|
|
input_text = "\n".join(input_lines) # No trailing newline
|
|
|
|
result = truncate_output(input_text, max_lines=5)
|
|
|
|
# Should handle truncation correctly without trailing newline
|
|
assert "[5 lines of output truncated]" in result
|
|
assert "Line 5" in result
|
|
assert "Line 9" in result
|
|
assert "Line 0" not in result
|
|
assert "Line 4" not in result
|