Files
X-Financial/server/tests/test_knowledge_tenant_security.py
caoxiaozhu 787bc3a481 feat(platform): close AI expense value loop
Add tenant-safe value, telemetry, connector, commercial, and production-readiness foundations.
2026-07-17 14:14:08 +08:00

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)