rp/tests/test_context.py

from rp.config import RECENT_MESSAGES_TO_KEEP
from rp.core.context import compress_context, should_compress_context


class TestContextManagement:

    def test_should_compress_context_below_threshold(self):
        messages = [{"role": "user", "content": "test"}] * 10
        assert should_compress_context(messages) is False

    def test_should_compress_context_above_threshold(self):
        messages = [{"role": "user", "content": "test"}] * 35
        assert should_compress_context(messages) is True

    def test_compress_context_preserves_system_message(self):
        messages = [
            {"role": "system", "content": "System prompt"},
            {"role": "user", "content": "Hello"},
            {"role": "assistant", "content": "Hi"},
        ] * 40  # Ensure compression
        compressed = compress_context(messages)
        assert compressed[0]["role"] == "system"
        assert "System prompt" in compressed[0]["content"]

    def test_compress_context_keeps_recent_messages(self):
        messages = [{"role": "user", "content": f"msg{i}"} for i in range(40)]
        compressed = compress_context(messages)
        # Should keep recent messages
        recent = compressed[-RECENT_MESSAGES_TO_KEEP:]
        assert len(recent) == RECENT_MESSAGES_TO_KEEP
        # Check that the messages are the most recent ones
        for i, msg in enumerate(recent):
            expected_index = 40 - RECENT_MESSAGES_TO_KEEP + i
            assert msg["content"] == f"msg{expected_index}"