replypier/backend/tests/knowledge/test_rechunk_knowledge.py

"""
Copyright 2024-2026 ReplyPier

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

from types import SimpleNamespace
from unittest.mock import MagicMock

import pytest

from scripts import rechunk_knowledge as script


def _row(id: str, content: str, **meta):
    return SimpleNamespace(
        id=id,
        name="https://x.com",
        content=content,
        meta_data={"url": id, **meta},
        filters={"name": "https://x.com", "agent_id": ["a1"], "org_id": "o1"},
    )


@pytest.fixture
def conn():
    return MagicMock()


@pytest.fixture
def embedder():
    embedder = MagicMock()
    embedder.get_embedding_and_usage.return_value = ([0.5, 0.5], None)
    return embedder


def _execute_returning(conn, oversized, taken=()):
    """First query lists oversized rows, second checks planned ids."""
    results = [MagicMock(fetchall=MagicMock(return_value=oversized)),
               MagicMock(fetchall=MagicMock(return_value=[SimpleNamespace(id=t) for t in taken]))]
    results += [MagicMock()] * 10
    conn.execute.side_effect = results


class TestProcessTable:
    def test_splits_row_and_replaces_it_then_commits(self, conn, embedder):
        big = _row("https://x.com/pricing", "para\n\n" * 400)
        _execute_returning(conn, [big])

        report = script.process_table(conn, "d_o1", embedder, size=300, dry_run=False)

        assert report.rows == 1 and report.failed == 0 and report.chunks > 1
        # Commit-as-you-go: never conn.begin() (the preceding SELECTs already
        # opened the transaction), one commit after insert + delete.
        conn.begin.assert_not_called()
        conn.commit.assert_called_once()
        insert_call, delete_call = conn.execute.call_args_list[2], conn.execute.call_args_list[3]
        records = insert_call.args[1]
        assert len(records) == report.chunks
        assert records[0]["id"] == "https://x.com/pricing::1"
        assert '"agent_id": ["a1"]' in records[0]["filters"]
        assert '"page_id": "https://x.com/pricing"' in records[0]["meta_data"]
        assert delete_call.args[1] == {"id": "https://x.com/pricing"}
        assert embedder.get_embedding_and_usage.call_count == report.chunks

    def test_dry_run_writes_nothing(self, conn, embedder):
        _execute_returning(conn, [_row("https://x.com/p", "x " * 400)])
        report = script.process_table(conn, "d_o1", None, size=300, dry_run=True)
        assert report.rows == 1 and report.chunks > 1
        conn.commit.assert_not_called()
        assert conn.execute.call_count == 2  # list + collision check only

    def test_collision_leaves_row_untouched(self, conn, embedder):
        _execute_returning(conn, [_row("https://x.com/p", "x " * 400)], taken=["https://x.com/p::2"])
        report = script.process_table(conn, "d_o1", embedder, size=300, dry_run=False)
        assert report.failed == 1 and report.collisions == ["https://x.com/p::2"]
        conn.commit.assert_not_called()

    def test_embedding_failure_keeps_original(self, conn, embedder):
        _execute_returning(conn, [_row("https://x.com/p", "x " * 400)])
        embedder.get_embedding_and_usage.side_effect = RuntimeError("model down")
        report = script.process_table(conn, "d_o1", embedder, size=300, dry_run=False)
        assert report.failed == 1
        conn.commit.assert_not_called()

    def test_write_failure_rolls_back(self, conn, embedder):
        _execute_returning(conn, [_row("https://x.com/p", "x " * 400)])
        conn.execute.side_effect = list(conn.execute.side_effect)[:2] + [RuntimeError("disk full")]
        report = script.process_table(conn, "d_o1", embedder, size=300, dry_run=False)
        assert report.failed == 1
        conn.rollback.assert_called_once()
        conn.commit.assert_not_called()

    def test_row_that_fits_one_chunk_is_skipped(self, conn, embedder):
        _execute_returning(conn, [_row("https://x.com/p", "short")])
        report = script.process_table(conn, "d_o1", embedder, size=300, dry_run=False)
        assert report.skipped == 1 and report.rows == 0

    def test_legacy_pdf_chunks_are_left_alone(self, conn, embedder):
        pdf = SimpleNamespace(id="manual_3", name="manual", content="p " * 400,
                              meta_data={"page": 3}, filters=None)
        _execute_returning(conn, [pdf])
        report = script.process_table(conn, "d_o1", embedder, size=300, dry_run=False)
        assert report.legacy_pdf == 1 and report.rows == 0
        conn.commit.assert_not_called()


class TestRenameLegacySuffixedPages:
    def test_renames_to_explicit_chunk_id(self, conn):
        conn.execute.return_value.fetchall.return_value = [
            SimpleNamespace(id="https://x.com/post_2", meta_data={"url": "https://x.com/post_2"})
        ]
        renamed = script.rename_legacy_suffixed_pages(conn, "d_o1", size=300, dry_run=False)
        assert renamed == 1
        params = conn.execute.call_args.args[1]
        assert params["new_id"] == "https://x.com/post_2::1"
        assert params["old_id"] == "https://x.com/post_2"
        assert '"page_id": "https://x.com/post_2"' in params["meta"]
        conn.commit.assert_called_once()

    def test_dry_run_only_reports(self, conn):
        conn.execute.return_value.fetchall.return_value = [
            SimpleNamespace(id="https://x.com/post_2", meta_data={})
        ]
        assert script.rename_legacy_suffixed_pages(conn, "d_o1", size=300, dry_run=True) == 1
        assert conn.execute.call_count == 1
        conn.commit.assert_not_called()


class TestVectorTables:
    def test_org_scopes_to_one_table(self, conn):
        assert script.vector_tables(conn, "abc") == ["d_abc"]
        conn.execute.assert_not_called()

    def test_lists_all_vector_tables(self, conn):
        conn.execute.return_value.fetchall.return_value = [
            SimpleNamespace(table_name="d_1"), SimpleNamespace(table_name="d_2")
        ]
        assert script.vector_tables(conn, None) == ["d_1", "d_2"]