govoplan-search/tests/test_postgres_search.py

from __future__ import annotations

import os
from types import SimpleNamespace
import unittest
import uuid

from sqlalchemy import create_engine, delete, inspect, select
from sqlalchemy.orm import Session

from govoplan_core.core.search import SearchDocument, SearchQuery
from govoplan_search.backend.db.models import (
    SearchIndexAclToken,
    SearchIndexChangeQueue,
    SearchIndexDocument,
    SearchIndexState,
)
from govoplan_search.backend.service import SearchIndexService


DATABASE_URL = os.environ.get(
    "GOVOPLAN_SEARCH_POSTGRES_URL",
    "",
).strip()


class _Registry:
    def manifests(self):
        return (
            SimpleNamespace(id="search"),
            SimpleNamespace(id="cases"),
        )

    def search_sources(self):
        return ()


@unittest.skipUnless(
    DATABASE_URL.startswith("postgresql"),
    "GOVOPLAN_SEARCH_POSTGRES_URL is not configured",
)
class PostgreSqlSearchTests(unittest.TestCase):
    @classmethod
    def setUpClass(cls) -> None:
        cls.engine = create_engine(DATABASE_URL)
        table_names = set(inspect(cls.engine).get_table_names())
        required = {
            SearchIndexDocument.__tablename__,
            SearchIndexAclToken.__tablename__,
            SearchIndexState.__tablename__,
            SearchIndexChangeQueue.__tablename__,
        }
        if not required.issubset(table_names):
            raise RuntimeError(
                "Search migrations were not applied before the "
                "PostgreSQL integration test."
            )

    @classmethod
    def tearDownClass(cls) -> None:
        cls.engine.dispose()

    def setUp(self) -> None:
        self.tenant_id = str(uuid.uuid4())
        self.other_tenant_id = str(uuid.uuid4())
        self.session = Session(self.engine)
        self.service = SearchIndexService(_Registry())
        self.principal = SimpleNamespace(
            tenant_id=self.tenant_id,
            account_id="account-allowed",
            membership_id="membership-allowed",
            identity_id=None,
            group_ids=frozenset(),
            role_ids=frozenset(),
            function_assignment_ids=frozenset(),
            scopes=frozenset({"cases:case:read"}),
        )

    def tearDown(self) -> None:
        document_ids = select(SearchIndexDocument.id).where(
            SearchIndexDocument.tenant_id.in_(
                (self.tenant_id, self.other_tenant_id)
            )
        )
        self.session.execute(
            delete(SearchIndexAclToken).where(
                SearchIndexAclToken.document_id.in_(document_ids)
            )
        )
        self.session.execute(
            delete(SearchIndexChangeQueue).where(
                SearchIndexChangeQueue.tenant_id.in_(
                    (self.tenant_id, self.other_tenant_id)
                )
            )
        )
        self.session.execute(
            delete(SearchIndexState).where(
                SearchIndexState.tenant_id.in_(
                    (self.tenant_id, self.other_tenant_id)
                )
            )
        )
        self.session.execute(
            delete(SearchIndexDocument).where(
                SearchIndexDocument.tenant_id.in_(
                    (self.tenant_id, self.other_tenant_id)
                )
            )
        )
        self.session.commit()
        self.session.close()

    def test_postgres_full_text_query_is_tenant_and_acl_bounded(
        self,
    ) -> None:
        self.service.upsert_document(
            self.session,
            self.principal,
            document=SearchDocument(
                tenant_id=self.tenant_id,
                module_id="cases",
                provider_id="cases.records",
                resource_type="case",
                resource_id="case-visible",
                title="Monthly parking permit review",
                url="/cases/case-visible",
                body="Compare the permit data with the payment register.",
                visibility="tenant",
            ),
        )
        self.service.upsert_document(
            self.session,
            self.principal,
            document=SearchDocument(
                tenant_id=self.tenant_id,
                module_id="cases",
                provider_id="cases.records",
                resource_type="case",
                resource_id="case-denied",
                title="Restricted parking permit",
                url="/cases/case-denied",
                visibility="restricted",
                acl_tokens=("account:someone-else",),
            ),
        )
        other_principal = SimpleNamespace(
            **{
                **vars(self.principal),
                "tenant_id": self.other_tenant_id,
            }
        )
        self.service.upsert_document(
            self.session,
            other_principal,
            document=SearchDocument(
                tenant_id=self.other_tenant_id,
                module_id="cases",
                provider_id="cases.records",
                resource_type="case",
                resource_id="case-other-tenant",
                title="Other tenant parking permit",
                url="/cases/case-other-tenant",
                visibility="tenant",
            ),
        )
        self.session.commit()

        page = self.service.search_page(
            self.session,
            self.principal,
            query=SearchQuery(
                text="parking permit",
                tenant_id=self.tenant_id,
            ),
        )

        self.assertEqual(
            ("case-visible",),
            tuple(item.resource_id for item in page.results),
        )
        self.assertEqual(
            "postgresql",
            self.service.diagnostics(
                self.session,
                self.principal,
            )["backend"],
        )


if __name__ == "__main__":
    unittest.main()