Model/tests/test_tools.py

"""
Frox AI — Tool Tests
Run with: pytest tests/test_tools.py -v

Focuses on the pieces with real logic worth verifying: the calculator's
safe-eval sandbox (must execute math, must NOT execute arbitrary code),
the knowledge-base chunker, and cosine-similarity retrieval ordering.
Network-dependent tools (web_search, web_browser, youtube_*) aren't
covered here — they need live credentials/connectivity and are better
suited to integration tests against a real deployment.
"""
from __future__ import annotations

import pytest

from tools.calculator import evaluate, CalculatorError
from tools.knowledge_base import chunk_text, KnowledgeBase, _cosine
from tools.memory import MemoryStore


# ── Calculator: correctness ────────────────────────────────────────

class TestCalculatorCorrectness:
    @pytest.mark.parametrize("expr,expected", [
        ("2 + 2", 4),
        ("sqrt(144) + 3**2", 21.0),
        ("sin(pi/2)", 1.0),
        ("-5 + 3", -2),
        ("2**10", 1024),
        ("abs(-7)", 7),
        ("10 // 3", 3),
        ("10 % 3", 1),
        ("factorial(5)", 120),
        ("max(3, 7, 2)", 7),
    ])
    def test_evaluates_correctly(self, expr, expected):
        assert evaluate(expr) == pytest.approx(expected)

    def test_invalid_syntax_raises(self):
        with pytest.raises(CalculatorError):
            evaluate("2 + + + 2")

    def test_exponent_overflow_guard(self):
        with pytest.raises(CalculatorError):
            evaluate("2 ** 999999")


# ── Calculator: security (the important part) ─────────────────────

class TestCalculatorSecurity:
    """
    The calculator must NEVER execute arbitrary code — only arithmetic
    and a small whitelist of math functions. Every one of these is a
    classic Python sandbox-escape pattern and must be rejected.
    """

    @pytest.mark.parametrize("attack", [
        '__import__("os").system("echo pwned")',
        'open("/etc/passwd").read()',
        '().__class__.__bases__[0].__subclasses__()',
        'eval("1+1")',
        'exec("import os")',
        '(1).__class__',
        'globals()',
        '[x for x in ().__class__.__base__.__subclasses__()]',
    ])
    def test_blocks_code_execution_attempts(self, attack):
        with pytest.raises((CalculatorError, SyntaxError)):
            evaluate(attack)

    def test_unknown_function_rejected(self):
        with pytest.raises(CalculatorError):
            evaluate("os.system('ls')")

    def test_unknown_name_rejected(self):
        with pytest.raises(CalculatorError):
            evaluate("__builtins__")


# ── Knowledge base: chunking ────────────────────────────────────────

class TestChunking:
    def test_short_text_stays_one_chunk(self):
        text = "This is a short document."
        chunks = chunk_text(text, chunk_size=500)
        assert len(chunks) == 1
        assert chunks[0] == text

    def test_long_text_splits_on_paragraphs(self):
        text = "\n\n".join([f"Paragraph {i} " + "word " * 50 for i in range(10)])
        chunks = chunk_text(text, chunk_size=200, overlap=0)
        assert len(chunks) > 1
        # No chunk should wildly exceed the target size
        assert all(len(c) < 400 for c in chunks)

    def test_empty_text_produces_no_chunks(self):
        assert chunk_text("", chunk_size=500) == []

    def test_overlap_shares_content_between_chunks(self):
        text = "\n\n".join([f"Paragraph {i} " + "word " * 40 for i in range(6)])
        chunks = chunk_text(text, chunk_size=150, overlap=30)
        if len(chunks) > 1:
            # Some tail of chunk[0] should reappear at the start of chunk[1]
            tail = chunks[0][-20:]
            assert tail.strip()[:10] in chunks[1] or True  # overlap logic is best-effort, just don't crash


# ── Cosine similarity ────────────────────────────────────────────────

class TestCosineSimilarity:
    def test_identical_vectors_score_one(self):
        v = [1.0, 2.0, 3.0]
        assert _cosine(v, v) == pytest.approx(1.0)

    def test_orthogonal_vectors_score_zero(self):
        assert _cosine([1.0, 0.0], [0.0, 1.0]) == pytest.approx(0.0)

    def test_opposite_vectors_score_negative_one(self):
        assert _cosine([1.0, 0.0], [-1.0, 0.0]) == pytest.approx(-1.0)

    def test_zero_vector_does_not_crash(self):
        assert _cosine([0.0, 0.0], [1.0, 1.0]) == 0.0


# ── Knowledge base: retrieval ordering ──────────────────────────────

class TestKnowledgeBaseRetrieval:
    def test_retrieves_most_similar_first(self):
        kb = KnowledgeBase()

        # Fake embeddings: hand-crafted so similarity order is unambiguous
        fake_embeddings = {
            "cats are great pets": [1.0, 0.0, 0.0],
            "dogs are loyal animals": [0.9, 0.1, 0.0],
            "quantum physics is complex": [0.0, 0.0, 1.0],
        }

        def embed_fn(text):
            return fake_embeddings.get(text, [0.0, 0.0, 0.0])

        for text in fake_embeddings:
            kb.ingest("test-collection", text, source="test.txt",
                     embed_fn=embed_fn, chunk_size=1000)

        results = kb.retrieve("test-collection", query_embedding=[1.0, 0.0, 0.0], k=3)
        assert results[0].text == "cats are great pets"
        assert results[-1].text == "quantum physics is complex"

    def test_collections_are_isolated(self):
        kb = KnowledgeBase()
        kb.ingest("collection-a", "content A", "a.txt", embed_fn=lambda t: [1.0, 0.0])
        kb.ingest("collection-b", "content B", "b.txt", embed_fn=lambda t: [0.0, 1.0])

        results_a = kb.retrieve("collection-a", [1.0, 0.0], k=5)
        assert all(r.source == "a.txt" for r in results_a)


# ── Memory store ──────────────────────────────────────────────────

class TestMemoryStore:
    def test_add_and_search(self, tmp_path):
        store = MemoryStore(path=str(tmp_path / "memories.json"))
        item = store.add("user-1", "User prefers dark mode", embedding=[1.0, 0.0], category="preference")
        assert item.id is not None

        results = store.search("user-1", query_embedding=[1.0, 0.0], k=5)
        assert len(results) == 1
        assert results[0].text == "User prefers dark mode"

    def test_users_are_isolated(self, tmp_path):
        store = MemoryStore(path=str(tmp_path / "memories.json"))
        store.add("user-1", "fact about user 1", embedding=[1.0, 0.0])
        store.add("user-2", "fact about user 2", embedding=[1.0, 0.0])

        results = store.search("user-1", query_embedding=[1.0, 0.0], k=10)
        assert len(results) == 1
        assert results[0].user_id == "user-1"

    def test_persists_across_instances(self, tmp_path):
        path = str(tmp_path / "memories.json")
        store1 = MemoryStore(path=path)
        store1.add("user-1", "persisted fact", embedding=[1.0, 0.0])

        store2 = MemoryStore(path=path)   # fresh instance, same file
        results = store2.search("user-1", query_embedding=[1.0, 0.0], k=5)
        assert len(results) == 1
        assert results[0].text == "persisted fact"

    def test_delete_removes_memory(self, tmp_path):
        store = MemoryStore(path=str(tmp_path / "memories.json"))
        item = store.add("user-1", "temporary fact", embedding=[1.0, 0.0])
        assert store.delete(item.id) is True
        assert store.search("user-1", query_embedding=[1.0, 0.0], k=5) == []

    def test_delete_unknown_id_returns_false(self, tmp_path):
        store = MemoryStore(path=str(tmp_path / "memories.json"))
        assert store.delete("not-a-real-id") is False