DataOps-platform/tests/integration/test_knowledge_dynamic_update.py

from __future__ import annotations

import os
import uuid

import pytest
from sqlalchemy import create_engine, text
from sqlalchemy.orm import Session

from app.core.knowledge.point_builder import build_knowledge_snapshot

pytestmark = pytest.mark.integration


class FixedEmbeddingProvider:
    profile_key = "qwen:test:1024"

    def __init__(self):
        self.calls: list[str] = []

    def embed(self, texts: list[str]) -> list[list[float]]:
        self.calls.extend(texts)
        return [[float(index + 1)] * 1024 for index, _text in enumerate(texts)]


@pytest.fixture()
def database_url():
    value = os.environ.get("TEST_DATABASE_URL")
    if not value:
        pytest.skip("TEST_DATABASE_URL is not configured")
    return value


def _snapshot(source_uid: str, domain_uid: str, version: int, owner: str):
    return build_knowledge_snapshot(
        "DataFlow",
        {
            "uid": source_uid,
            "version": version,
            "name": "customer_sync",
            "purpose": "sync customer data",
            "owner": owner,
            "business_domain_uid": domain_uid,
        },
    )


def test_sql_repository_atomically_switches_active_points_and_reuses_embeddings(
    database_url,
):
    from app.core.knowledge.repository import SqlKnowledgeRepository
    from app.core.knowledge.sync import KnowledgeSyncService

    engine = create_engine(database_url)
    source_uid = str(uuid.uuid4())
    domain_uid = str(uuid.uuid4())
    profile_id = str(uuid.uuid4())
    try:
        with engine.begin() as connection:
            connection.execute(
                text(
                    """
                    INSERT INTO public.knowledge_embedding_profiles (
                        id, provider, model, dimension, distance, status, config_hash,
                        activated_at
                    ) VALUES (
                        CAST(:id AS uuid), 'qwen', 'test', 1024, 'cosine',
                        'active', :config_hash, CURRENT_TIMESTAMP
                    )
                    """
                ),
                {"id": profile_id, "config_hash": "a" * 64},
            )

        first_embedder = FixedEmbeddingProvider()
        with Session(engine) as session, session.begin():
            repository = SqlKnowledgeRepository(
                session,
                embedding_profile_id=profile_id,
                embedding_profile_key=first_embedder.profile_key,
                generation=1,
                workspace=f"dataops-global-{domain_uid}-g1",
            )
            first = KnowledgeSyncService(
                repository=repository, embedder=first_embedder
            ).sync(_snapshot(source_uid, domain_uid, 1, "team-a"))
        assert first.embedded_chunk_count == 3

        second_embedder = FixedEmbeddingProvider()
        with Session(engine) as session, session.begin():
            repository = SqlKnowledgeRepository(
                session,
                embedding_profile_id=profile_id,
                embedding_profile_key=second_embedder.profile_key,
                generation=1,
                workspace=f"dataops-global-{domain_uid}-g1",
            )
            second = KnowledgeSyncService(
                repository=repository, embedder=second_embedder
            ).sync(_snapshot(source_uid, domain_uid, 2, "team-b"))

        assert second.diff_counts == {
            "added": 0,
            "modified": 1,
            "deleted": 0,
            "unchanged": 2,
        }
        assert second_embedder.calls == ["team-b"]

        with engine.connect() as connection:
            active_document = connection.execute(
                text(
                    "SELECT object_version, point_set_hash FROM public.governance_documents "
                    "WHERE object_uid = CAST(:uid AS uuid) AND status = 'active'"
                ),
                {"uid": source_uid},
            ).one()
            active_points = connection.execute(
                text(
                    "SELECT semantic_path, content FROM public.knowledge_points "
                    "WHERE source_uid = CAST(:uid AS uuid) AND status = 'active' "
                    "ORDER BY semantic_path"
                ),
                {"uid": source_uid},
            ).all()
            change_sets = connection.execute(
                text(
                    "SELECT id::text, status FROM public.knowledge_change_sets "
                    "WHERE source_uid = CAST(:uid AS uuid) ORDER BY source_revision"
                ),
                {"uid": source_uid},
            ).all()

        assert active_document[0] == 2
        assert active_document[1]
        assert dict(active_points)["owner"] == "team-b"
        assert [row[1] for row in change_sets] == [
            "canonical_active",
            "canonical_active",
        ]

        from app.core.knowledge.admin import rollback_change_set

        with Session(engine) as session, session.begin():
            assert rollback_change_set(session, change_sets[1][0]) is True
        with engine.connect() as connection:
            rolled_back_document = connection.execute(
                text(
                    "SELECT object_version FROM public.governance_documents "
                    "WHERE object_uid = CAST(:uid AS uuid) AND status = 'active'"
                ),
                {"uid": source_uid},
            ).scalar_one()
            rolled_back_owner = connection.execute(
                text(
                    "SELECT content FROM public.knowledge_points "
                    "WHERE point_key = :point_key AND status = 'active'"
                ),
                {"point_key": f"DataFlow/{source_uid}/owner"},
            ).scalar_one()
        assert rolled_back_document == 1
        assert rolled_back_owner == "team-a"
    finally:
        with engine.begin() as connection:
            connection.execute(
                text(
                    "DELETE FROM public.knowledge_index_projections "
                    "WHERE document_id IN (SELECT id FROM public.governance_documents "
                    "WHERE object_uid = CAST(:uid AS uuid))"
                ),
                {"uid": source_uid},
            )
            connection.execute(
                text(
                    "DELETE FROM public.knowledge_chunk_embeddings "
                    "WHERE chunk_id IN (SELECT c.id FROM public.governance_chunks c "
                    "JOIN public.governance_documents d ON d.id = c.document_id "
                    "WHERE d.object_uid = CAST(:uid AS uuid))"
                ),
                {"uid": source_uid},
            )
            connection.execute(
                text(
                    "DELETE FROM public.governance_documents "
                    "WHERE object_uid = CAST(:uid AS uuid)"
                ),
                {"uid": source_uid},
            )
            connection.execute(
                text(
                    "DELETE FROM public.knowledge_point_dependencies "
                    "WHERE from_point_key LIKE :prefix"
                ),
                {"prefix": f"DataFlow/{source_uid}/%"},
            )
            connection.execute(
                text(
                    "DELETE FROM public.knowledge_points "
                    "WHERE source_uid = CAST(:uid AS uuid)"
                ),
                {"uid": source_uid},
            )
            connection.execute(
                text(
                    "DELETE FROM public.knowledge_change_sets "
                    "WHERE source_uid = CAST(:uid AS uuid)"
                ),
                {"uid": source_uid},
            )
            connection.execute(
                text(
                    "DELETE FROM public.knowledge_embedding_profiles "
                    "WHERE id = CAST(:id AS uuid)"
                ),
                {"id": profile_id},
            )
        engine.dispose()