231 lines
8.9 KiB
Python
231 lines
8.9 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import json
|
||
|
|
|
||
|
|
import pytest
|
||
|
|
from sqlalchemy import create_engine
|
||
|
|
from sqlalchemy.orm import sessionmaker
|
||
|
|
from sqlalchemy.pool import StaticPool
|
||
|
|
|
||
|
|
from app.api.deps import CurrentUserContext
|
||
|
|
from app.db.base import Base
|
||
|
|
from app.models.agent_run import AgentRun
|
||
|
|
from app.models.tenant import Tenant
|
||
|
|
from app.services.knowledge import KnowledgeService
|
||
|
|
from app.services.knowledge_rag import KnowledgeRagService
|
||
|
|
from app.services.knowledge_run_scope import resolve_trusted_knowledge_run_tenant
|
||
|
|
from app.services.knowledge_scheduler import KnowledgeIndexScheduler
|
||
|
|
from app.services.knowledge_tenant_scope import PLATFORM_KNOWLEDGE_SCOPE
|
||
|
|
|
||
|
|
|
||
|
|
def _user(tenant_id: str) -> CurrentUserContext:
|
||
|
|
return CurrentUserContext(
|
||
|
|
username=f"admin-{tenant_id}",
|
||
|
|
name=f"管理员 {tenant_id}",
|
||
|
|
role_codes=["manager"],
|
||
|
|
is_admin=True,
|
||
|
|
tenant_id=tenant_id,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def test_tenant_document_metadata_paths_and_reads_are_isolated(tmp_path, monkeypatch) -> None:
|
||
|
|
monkeypatch.setattr(
|
||
|
|
KnowledgeRagService,
|
||
|
|
"get_document_status_map",
|
||
|
|
lambda _self, _document_ids: {},
|
||
|
|
)
|
||
|
|
monkeypatch.setattr(KnowledgeRagService, "delete_document", lambda *_args: None)
|
||
|
|
tenant_a = KnowledgeService(storage_root=tmp_path, tenant_id="tenant-a")
|
||
|
|
tenant_b = KnowledgeService(storage_root=tmp_path, tenant_id="tenant-b")
|
||
|
|
|
||
|
|
doc_a = tenant_a.upload_document("报销制度", "制度.txt", b"tenant-a", _user("tenant-a"))
|
||
|
|
doc_b = tenant_b.upload_document("报销制度", "制度.txt", b"tenant-b", _user("tenant-b"))
|
||
|
|
|
||
|
|
assert doc_a.id != doc_b.id
|
||
|
|
path_a, _, _ = tenant_a.get_document_content(doc_a.id)
|
||
|
|
path_b, _, _ = tenant_b.get_document_content(doc_b.id)
|
||
|
|
assert path_a.read_bytes() == b"tenant-a"
|
||
|
|
assert path_b.read_bytes() == b"tenant-b"
|
||
|
|
assert path_a.is_relative_to(tmp_path / "knowledge" / "tenants" / "tenant-a")
|
||
|
|
assert path_b.is_relative_to(tmp_path / "knowledge" / "tenants" / "tenant-b")
|
||
|
|
with pytest.raises(FileNotFoundError):
|
||
|
|
tenant_a.get_document_content(doc_b.id)
|
||
|
|
with pytest.raises(FileNotFoundError):
|
||
|
|
tenant_b.get_document_detail(doc_a.id)
|
||
|
|
|
||
|
|
|
||
|
|
def test_platform_documents_are_explicitly_read_only_and_visible_to_tenants(
|
||
|
|
tmp_path,
|
||
|
|
monkeypatch,
|
||
|
|
) -> None:
|
||
|
|
legacy_folder = tmp_path / "knowledge" / "制度政策"
|
||
|
|
legacy_folder.mkdir(parents=True)
|
||
|
|
legacy_file = legacy_folder / "platform-doc__平台制度.txt"
|
||
|
|
legacy_file.write_bytes(b"platform-policy")
|
||
|
|
(tmp_path / "knowledge" / ".index.json").write_text(
|
||
|
|
json.dumps(
|
||
|
|
{
|
||
|
|
"version": 1,
|
||
|
|
"documents": [
|
||
|
|
{
|
||
|
|
"id": "platform-doc",
|
||
|
|
"folder": "制度政策",
|
||
|
|
"original_name": "平台制度.txt",
|
||
|
|
"stored_name": legacy_file.name,
|
||
|
|
"mime_type": "text/plain",
|
||
|
|
"extension": "txt",
|
||
|
|
"size_bytes": len(b"platform-policy"),
|
||
|
|
"sha256": "",
|
||
|
|
"created_at": "2026-07-17T00:00:00+00:00",
|
||
|
|
"updated_at": "2026-07-17T00:00:00+00:00",
|
||
|
|
"uploaded_by": "平台",
|
||
|
|
"version_number": 1,
|
||
|
|
"ingest_status": 1,
|
||
|
|
}
|
||
|
|
],
|
||
|
|
},
|
||
|
|
ensure_ascii=False,
|
||
|
|
),
|
||
|
|
encoding="utf-8",
|
||
|
|
)
|
||
|
|
monkeypatch.setattr(
|
||
|
|
KnowledgeRagService,
|
||
|
|
"get_document_status_map",
|
||
|
|
lambda _self, _document_ids: {},
|
||
|
|
)
|
||
|
|
platform = KnowledgeService(storage_root=tmp_path, scope=PLATFORM_KNOWLEDGE_SCOPE)
|
||
|
|
platform.ensure_library_ready()
|
||
|
|
tenant = KnowledgeService(storage_root=tmp_path, tenant_id="tenant-a")
|
||
|
|
|
||
|
|
listed = {item.id: item for item in tenant.list_library().documents}
|
||
|
|
assert listed["platform-doc"].scope == PLATFORM_KNOWLEDGE_SCOPE
|
||
|
|
assert listed["platform-doc"].readOnly is True
|
||
|
|
assert tenant.get_document_content("platform-doc")[0].read_bytes() == b"platform-policy"
|
||
|
|
with pytest.raises(FileNotFoundError):
|
||
|
|
tenant.delete_document("platform-doc")
|
||
|
|
with pytest.raises(ValueError, match="只读"):
|
||
|
|
platform.upload_document("制度政策", "new.txt", b"x", _user("tenant-a"))
|
||
|
|
|
||
|
|
|
||
|
|
def test_unscoped_knowledge_services_fail_closed(tmp_path) -> None:
|
||
|
|
with pytest.raises(ValueError, match="显式提供"):
|
||
|
|
KnowledgeService(storage_root=tmp_path)
|
||
|
|
with pytest.raises(ValueError, match="显式提供"):
|
||
|
|
KnowledgeRagService(storage_root=tmp_path)
|
||
|
|
|
||
|
|
|
||
|
|
def test_lightrag_workspace_cache_and_local_chunks_are_tenant_namespaced(tmp_path) -> None:
|
||
|
|
service_a = KnowledgeRagService(storage_root=tmp_path, tenant_id="tenant-a")
|
||
|
|
service_b = KnowledgeRagService(storage_root=tmp_path, tenant_id="tenant-b")
|
||
|
|
assert service_a.storage_scope.workspace != service_b.storage_scope.workspace
|
||
|
|
assert service_a.storage_scope.runtime_cache_key != service_b.storage_scope.runtime_cache_key
|
||
|
|
assert "tenant-a" not in service_a.storage_scope.workspace
|
||
|
|
|
||
|
|
for service, marker in ((service_a, "A 租户专属限额"), (service_b, "B 租户专属限额")):
|
||
|
|
workspace = service.storage_scope.lightrag_root / service.storage_scope.workspace
|
||
|
|
workspace.mkdir(parents=True)
|
||
|
|
(workspace / "kv_store_text_chunks.json").write_text(
|
||
|
|
json.dumps(
|
||
|
|
{
|
||
|
|
f"chunk-{marker[0]}": {
|
||
|
|
"_id": f"chunk-{marker[0]}",
|
||
|
|
"full_doc_id": f"doc-{marker[0]}",
|
||
|
|
"chunk_order_index": 1,
|
||
|
|
"file_path": f"/tmp/doc-{marker[0]}__制度.txt",
|
||
|
|
"content": f"报销限额规定：{marker}，提交前必须校验。",
|
||
|
|
}
|
||
|
|
},
|
||
|
|
ensure_ascii=False,
|
||
|
|
),
|
||
|
|
encoding="utf-8",
|
||
|
|
)
|
||
|
|
|
||
|
|
result_a = service_a.query_knowledge("A 租户专属限额是什么？", limit=2)
|
||
|
|
result_b = service_b.query_knowledge("B 租户专属限额是什么？", limit=2)
|
||
|
|
assert "A 租户专属限额" in result_a["hits"][0]["content"]
|
||
|
|
assert "B 租户专属限额" in result_b["hits"][0]["content"]
|
||
|
|
assert all("B 租户" not in item["content"] for item in result_a["hits"])
|
||
|
|
assert all("A 租户" not in item["content"] for item in result_b["hits"])
|
||
|
|
|
||
|
|
|
||
|
|
def test_scheduler_enumerates_only_active_tenant_registry_rows(monkeypatch) -> None:
|
||
|
|
assert "tenants" in Base.metadata.tables
|
||
|
|
engine = create_engine(
|
||
|
|
"sqlite+pysqlite:///:memory:",
|
||
|
|
connect_args={"check_same_thread": False},
|
||
|
|
poolclass=StaticPool,
|
||
|
|
)
|
||
|
|
Tenant.__table__.create(bind=engine)
|
||
|
|
factory = sessionmaker(bind=engine, autoflush=False, autocommit=False)
|
||
|
|
with factory() as db:
|
||
|
|
db.add_all(
|
||
|
|
[
|
||
|
|
Tenant(tenant_id="tenant-a", tenant_code="A", name="A", status="active"),
|
||
|
|
Tenant(
|
||
|
|
tenant_id="tenant-b",
|
||
|
|
tenant_code="B",
|
||
|
|
name="B",
|
||
|
|
status="suspended",
|
||
|
|
),
|
||
|
|
]
|
||
|
|
)
|
||
|
|
db.commit()
|
||
|
|
|
||
|
|
seen_tenants: list[str] = []
|
||
|
|
|
||
|
|
class FakeDispatch:
|
||
|
|
def __init__(self, _db) -> None:
|
||
|
|
pass
|
||
|
|
|
||
|
|
def queue_sync(self, *, current_user, **_kwargs):
|
||
|
|
seen_tenants.append(current_user.tenant_id)
|
||
|
|
return type(
|
||
|
|
"Result",
|
||
|
|
(),
|
||
|
|
{
|
||
|
|
"agent_run_id": "",
|
||
|
|
"document_ids": [],
|
||
|
|
"reused": False,
|
||
|
|
"summary": "no changes",
|
||
|
|
},
|
||
|
|
)()
|
||
|
|
|
||
|
|
monkeypatch.setattr(
|
||
|
|
"app.services.knowledge_scheduler.get_session_factory",
|
||
|
|
lambda: factory,
|
||
|
|
)
|
||
|
|
monkeypatch.setattr(
|
||
|
|
"app.services.knowledge_scheduler.KnowledgeSyncDispatchService",
|
||
|
|
FakeDispatch,
|
||
|
|
)
|
||
|
|
KnowledgeIndexScheduler()._run_incremental_sync()
|
||
|
|
assert seen_tenants == ["tenant-a"]
|
||
|
|
|
||
|
|
|
||
|
|
def test_index_worker_rederives_tenant_from_trusted_agent_run() -> None:
|
||
|
|
engine = create_engine(
|
||
|
|
"sqlite+pysqlite:///:memory:",
|
||
|
|
connect_args={"check_same_thread": False},
|
||
|
|
poolclass=StaticPool,
|
||
|
|
)
|
||
|
|
AgentRun.__table__.create(bind=engine)
|
||
|
|
factory = sessionmaker(bind=engine, autoflush=False, autocommit=False)
|
||
|
|
with factory() as db:
|
||
|
|
run = AgentRun(
|
||
|
|
run_id="run-tenant-bound",
|
||
|
|
agent="hermes",
|
||
|
|
source="schedule",
|
||
|
|
route_json={"tenant_id": "tenant-a"},
|
||
|
|
ontology_json={"tenant_id": "tenant-a"},
|
||
|
|
permission_level="read",
|
||
|
|
status="running",
|
||
|
|
)
|
||
|
|
db.add(run)
|
||
|
|
db.commit()
|
||
|
|
assert resolve_trusted_knowledge_run_tenant(db, run.run_id) == "tenant-a"
|
||
|
|
|
||
|
|
run.ontology_json = {"tenant_id": "tenant-b"}
|
||
|
|
db.commit()
|
||
|
|
with pytest.raises(ValueError, match="冲突"):
|
||
|
|
resolve_trusted_knowledge_run_tenant(db, run.run_id)
|