cognee 0.5.1.dev0__py3-none-any.whl → 0.5.2.dev0__py3-none-any.whl

cognee/tests/unit/modules/search/test_search.py

@@ -0,0 +1,176 @@
+import types
+from uuid import uuid4, uuid5, NAMESPACE_OID
+
+import pytest
+
+from cognee.modules.search.models.SearchResultPayload import SearchResultPayload
+from cognee.modules.search.types import SearchType
+
+
+def _make_user(user_id: str = "u1", tenant_id=None):
+    return types.SimpleNamespace(id=user_id, tenant_id=tenant_id)
+
+
+def _make_dataset(*, name="ds", tenant_id="t1", dataset_id=None, owner_id=None):
+    return types.SimpleNamespace(
+        id=uuid5(NAMESPACE_OID, name),
+        name=name,
+        tenant_id=uuid5(NAMESPACE_OID, tenant_id),
+        owner_id=owner_id or uuid4(),
+    )
+
+
+@pytest.fixture
+def search_mod():
+    import importlib
+
+    return importlib.import_module("cognee.modules.search.methods.search")
+
+
+@pytest.fixture(autouse=True)
+def _patch_side_effect_boundaries(monkeypatch, search_mod):
+    """
+    Keep production logic; patch only unavoidable side-effect boundaries.
+    """
+
+    async def dummy_log_query(_query_text, _query_type, _user_id):
+        return types.SimpleNamespace(id="qid-1")
+
+    async def dummy_log_result(*_args, **_kwargs):
+        return None
+
+    monkeypatch.setattr(search_mod, "send_telemetry", lambda *a, **k: None)
+    monkeypatch.setattr(search_mod, "log_query", dummy_log_query)
+    monkeypatch.setattr(search_mod, "log_result", dummy_log_result)
+
+    yield
+
+
+@pytest.mark.asyncio
+async def test_search_access_control_returns_dataset_shaped_dicts(monkeypatch, search_mod):
+    user = _make_user()
+    ds = _make_dataset(name="ds1", tenant_id="t1")
+
+    async def dummy_authorized_search(**_kwargs):
+        assert _kwargs["dataset_ids"] == [ds.id]
+        return [
+            SearchResultPayload(
+                result_object="object",
+                context=["ctx"],
+                completion=["r"],
+                search_type=SearchType.CHUNKS,
+                dataset_name=ds.name,
+                dataset_id=ds.id,
+                dataset_tenant_id=ds.tenant_id,
+            )
+        ]
+
+    monkeypatch.setattr(search_mod, "backend_access_control_enabled", lambda: True)
+    monkeypatch.setattr(search_mod, "authorized_search", dummy_authorized_search)
+
+    out_non_verbose = await search_mod.search(
+        query_text="q",
+        query_type=SearchType.CHUNKS,
+        dataset_ids=[ds.id],
+        user=user,
+        verbose=False,
+    )
+
+    assert out_non_verbose == [
+        {
+            "search_result": ["r"],
+            "dataset_id": ds.id,
+            "dataset_name": "ds1",
+            "dataset_tenant_id": uuid5(NAMESPACE_OID, "t1"),
+        }
+    ]
+
+
+@pytest.mark.asyncio
+async def test_search_access_control_only_context_returns_dataset_shaped_dicts(
+    monkeypatch, search_mod
+):
+    user = _make_user()
+    ds = _make_dataset(name="ds1", tenant_id="t1")
+
+    async def dummy_authorized_search(**_kwargs):
+        return [
+            SearchResultPayload(
+                result_object="object",
+                context=["ctx"],
+                completion=None,
+                search_type=SearchType.CHUNKS,
+                dataset_name=ds.name,
+                dataset_id=ds.id,
+                dataset_tenant_id=ds.tenant_id,
+            )
+        ]
+
+    monkeypatch.setattr(search_mod, "backend_access_control_enabled", lambda: True)
+    monkeypatch.setattr(search_mod, "authorized_search", dummy_authorized_search)
+
+    out = await search_mod.search(
+        query_text="q",
+        query_type=SearchType.CHUNKS,
+        dataset_ids=[ds.id],
+        user=user,
+        only_context=True,
+    )
+
+    assert out == [
+        {
+            "search_result": ["ctx"],
+            "dataset_id": ds.id,
+            "dataset_name": "ds1",
+            "dataset_tenant_id": uuid5(NAMESPACE_OID, "t1"),
+        }
+    ]
+
+
+@pytest.mark.asyncio
+async def test_authorized_search_delegates_to_search_in_datasets_context(monkeypatch, search_mod):
+    user = _make_user()
+    ds = _make_dataset(name="ds1")
+
+    async def dummy_get_authorized_existing_datasets(*_args, **_kwargs):
+        return [ds]
+
+    expected = [
+        SearchResultPayload(
+            result_object="object",
+            context="text",
+            completion="test",
+            search_type=SearchType.CHUNKS,
+            dataset_name=ds.name,
+            dataset_id=ds.id,
+            dataset_tenant_id=ds.tenant_id,
+        )
+    ]
+
+    async def dummy_search_in_datasets_context(**_kwargs):
+        return [
+            SearchResultPayload(
+                result_object="object",
+                context="text",
+                completion="test",
+                search_type=SearchType.CHUNKS,
+                dataset_name=ds.name,
+                dataset_id=ds.id,
+                dataset_tenant_id=ds.tenant_id,
+            )
+        ]
+
+    monkeypatch.setattr(
+        search_mod, "get_authorized_existing_datasets", dummy_get_authorized_existing_datasets
+    )
+    monkeypatch.setattr(search_mod, "search_in_datasets_context", dummy_search_in_datasets_context)
+
+    out = await search_mod.authorized_search(
+        query_type=SearchType.CHUNKS,
+        query_text="q",
+        user=user,
+        dataset_ids=[ds.id],
+        only_context=False,
+    )
+
+    assert out == expected

cognee/tests/unit/modules/search/test_search_prepare_search_result_contract.py

@@ -0,0 +1,190 @@
+## The Objective of these tests is to cover the search - prepare search results behavior (later to be removed)
+
+import types
+from uuid import uuid4, uuid5, UUID, NAMESPACE_OID
+
+import pytest
+from pydantic import BaseModel
+
+from cognee.modules.search.models.SearchResultPayload import SearchResultPayload
+from cognee.modules.search.types import SearchType
+
+
+class DummyDataset(BaseModel):
+    id: object
+    name: str
+    tenant_id: UUID | None = None
+    owner_id: object
+
+
+def _ds(name="ds1", tenant_id=uuid5(NAMESPACE_OID, "t1")):
+    return DummyDataset(
+        id=uuid5(NAMESPACE_OID, name), name=name, tenant_id=tenant_id, owner_id=uuid4()
+    )
+
+
+@pytest.fixture
+def search_mod():
+    import importlib
+
+    return importlib.import_module("cognee.modules.search.methods.search")
+
+
+@pytest.fixture(autouse=True)
+def _patch_search_side_effects(monkeypatch, search_mod):
+    """
+    These tests validate prepare_search_result behavior *through* search.py.
+    We only patch unavoidable side effects (telemetry + query/result logging).
+    """
+
+    async def dummy_log_query(_query_text, _query_type, _user_id):
+        return types.SimpleNamespace(id="qid-1")
+
+    async def dummy_log_result(*_args, **_kwargs):
+        return None
+
+    monkeypatch.setattr(search_mod, "send_telemetry", lambda *a, **k: None)
+    monkeypatch.setattr(search_mod, "log_query", dummy_log_query)
+    monkeypatch.setattr(search_mod, "log_result", dummy_log_result)
+
+    yield
+
+
+@pytest.mark.asyncio
+async def test_search_backend_access_verbose(monkeypatch, search_mod):
+    user = types.SimpleNamespace(id="u1", tenant_id=None)
+    ds = _ds("ds1")
+
+    async def dummy_authorized_search(**_kwargs):
+        return [
+            SearchResultPayload(
+                result_object="object",
+                context="text",
+                completion="test",
+                search_type=SearchType.CHUNKS,
+                dataset_name=ds.name,
+                dataset_id=ds.id,
+                dataset_tenant_id=ds.tenant_id,
+            )
+        ]
+
+    monkeypatch.setattr(search_mod, "backend_access_control_enabled", lambda: True)
+    monkeypatch.setattr(search_mod, "authorized_search", dummy_authorized_search)
+
+    out = await search_mod.search(
+        query_text="q",
+        query_type=SearchType.CHUNKS,
+        dataset_ids=[ds.id],
+        user=user,
+        verbose=True,
+    )
+
+    assert out[0]["dataset_name"] == "ds1"
+    assert out[0]["dataset_tenant_id"] == uuid5(NAMESPACE_OID, "t1")
+    assert out[0]["dataset_id"] == uuid5(NAMESPACE_OID, "ds1")
+    assert out[0]["objects_result"] == "object"
+    assert out[0]["context_result"] == "text"
+    assert out[0]["text_result"] == "test"
+
+
+@pytest.mark.asyncio
+async def test_search_no_backend_access_verbose(monkeypatch, search_mod):
+    user = types.SimpleNamespace(id="u1", tenant_id=None)
+    ds = _ds("ds1")
+
+    async def dummy_authorized_search(**_kwargs):
+        return [
+            SearchResultPayload(
+                result_object="object",
+                context="text",
+                completion="test",
+                search_type=SearchType.CHUNKS,
+                dataset_name=ds.name,
+                dataset_id=ds.id,
+                dataset_tenant_id=ds.tenant_id,
+            )
+        ]
+
+    monkeypatch.setattr(search_mod, "backend_access_control_enabled", lambda: False)
+    monkeypatch.setattr(search_mod, "authorized_search", dummy_authorized_search)
+
+    out = await search_mod.search(
+        query_text="q",
+        query_type=SearchType.CHUNKS,
+        dataset_ids=[ds.id],
+        user=user,
+        verbose=True,
+    )
+
+    assert out[0]["objects_result"] == "object"
+    assert out[0]["context_result"] == "text"
+    assert out[0]["text_result"] == "test"
+    assert out[0].get("dataset_name") is None
+
+
+@pytest.mark.asyncio
+async def test_search_backend_access(monkeypatch, search_mod):
+    user = types.SimpleNamespace(id="u1", tenant_id=None)
+    ds = _ds("ds1")
+
+    async def dummy_authorized_search(**_kwargs):
+        return [
+            SearchResultPayload(
+                result_object="object",
+                context="text",
+                completion="test",
+                search_type=SearchType.CHUNKS,
+                dataset_name=ds.name,
+                dataset_id=ds.id,
+                dataset_tenant_id=ds.tenant_id,
+            )
+        ]
+
+    monkeypatch.setattr(search_mod, "backend_access_control_enabled", lambda: True)
+    monkeypatch.setattr(search_mod, "authorized_search", dummy_authorized_search)
+
+    out = await search_mod.search(
+        query_text="q",
+        query_type=SearchType.CHUNKS,
+        dataset_ids=[ds.id],
+        user=user,
+    )
+
+    assert out[0]["dataset_name"] == "ds1"
+    assert out[0]["dataset_id"] == uuid5(NAMESPACE_OID, "ds1")
+    assert out[0]["dataset_tenant_id"] == uuid5(NAMESPACE_OID, "t1")
+    assert (
+        out[0]["search_result"] == "test"
+    )  # Search result should be resolved to completion if it exists
+
+
+@pytest.mark.asyncio
+async def test_search_no_backend_access(monkeypatch, search_mod):
+    user = types.SimpleNamespace(id="u1", tenant_id=None)
+    ds = _ds("ds1")
+
+    async def dummy_authorized_search(**_kwargs):
+        return [
+            SearchResultPayload(
+                result_object="object",
+                context="text",
+                completion="test",
+                search_type=SearchType.CHUNKS,
+                dataset_name=ds.name,
+                dataset_id=ds.id,
+                dataset_tenant_id=ds.tenant_id,
+            )
+        ]
+
+    monkeypatch.setattr(search_mod, "backend_access_control_enabled", lambda: False)
+    monkeypatch.setattr(search_mod, "authorized_search", dummy_authorized_search)
+
+    out = await search_mod.search(
+        query_text="q",
+        query_type=SearchType.CHUNKS,
+        dataset_ids=[ds.id],
+        user=user,
+    )
+
+    assert isinstance(out, list)
+    assert out[0] == "test"  # Search result should be resolved to completion if it exists