75 lines
3.8 KiB
Python
75 lines
3.8 KiB
Python
|
import unittest
|
||
|
|
from types import SimpleNamespace
|
||
|
|
|
||
|
|
from fastapi import HTTPException
|
||
|
|
from sqlmodel import SQLModel, Session, create_engine, select
|
||
|
|
|
||
|
|
from app.access import can_access, require_kb, require_memory
|
||
|
|
from app.main import (create_kb, grant_knowledge_base, grant_memory, list_kbs,
|
||
|
|
list_memories, list_admin_knowledge_bases, list_admin_memories)
|
||
|
|
from app.models import AppUser, KnowledgeBase, Memory
|
||
|
|
|
||
|
|
|
||
|
|
def request_for(user: AppUser):
|
||
|
|
return SimpleNamespace(state=SimpleNamespace(user=user))
|
||
|
|
|
||
|
|
|
||
|
|
class ResourceAccessTests(unittest.TestCase):
|
||
|
|
def setUp(self):
|
||
|
|
self.engine = create_engine("sqlite://")
|
||
|
|
SQLModel.metadata.create_all(self.engine)
|
||
|
|
self.db = Session(self.engine)
|
||
|
|
self.admin = AppUser(username="admin", display_name="管理员", role="admin")
|
||
|
|
self.user1 = AppUser(username="user_one", display_name="用户一", role="editor")
|
||
|
|
self.user2 = AppUser(username="user_two", display_name="用户二", role="editor")
|
||
|
|
self.db.add_all([self.admin, self.user1, self.user2])
|
||
|
|
self.db.commit()
|
||
|
|
for user in (self.admin, self.user1, self.user2):
|
||
|
|
self.db.refresh(user)
|
||
|
|
|
||
|
|
def tearDown(self):
|
||
|
|
self.db.close()
|
||
|
|
|
||
|
|
def test_scope_filters_private_global_and_shared_resources(self):
|
||
|
|
private = create_kb({"name": "用户一私有库", "visibility": "global"},
|
||
|
|
self.db, request_for(self.user1))
|
||
|
|
global_kb = create_kb({"name": "管理员全局库", "visibility": "global"},
|
||
|
|
self.db, request_for(self.admin))
|
||
|
|
shared = create_kb({"name": "授权库", "visibility": "private"},
|
||
|
|
self.db, request_for(self.admin))
|
||
|
|
grant_knowledge_base(shared["id"], {"user_id": self.user2.id, "permission": "use"}, self.db)
|
||
|
|
|
||
|
|
user1_ids = {row["id"] for row in list_kbs(self.db, request_for(self.user1))}
|
||
|
|
user2_ids = {row["id"] for row in list_kbs(self.db, request_for(self.user2))}
|
||
|
|
self.assertIn(private["id"], user1_ids)
|
||
|
|
self.assertNotIn(private["id"], user2_ids)
|
||
|
|
self.assertIn(global_kb["id"], user2_ids)
|
||
|
|
self.assertIn(shared["id"], user2_ids)
|
||
|
|
self.assertEqual(require_kb(self.db, shared["id"], request_for(self.user2), "use").id,
|
||
|
|
shared["id"])
|
||
|
|
with self.assertRaises(HTTPException):
|
||
|
|
require_kb(self.db, shared["id"], request_for(self.user1), "use")
|
||
|
|
|
||
|
|
def test_grant_permissions_are_hierarchical_and_admin_views_include_grants(self):
|
||
|
|
kb = KnowledgeBase(name="授权校验", organization_id=self.user1.organization_id,
|
||
|
|
owner_user_id=self.user1.id, visibility="shared")
|
||
|
|
memory = Memory(agent_id=0, content="共享记忆", organization_id=self.user1.organization_id,
|
||
|
|
owner_user_id=self.user1.id, visibility="shared")
|
||
|
|
self.db.add(kb); self.db.add(memory); self.db.commit()
|
||
|
|
self.db.refresh(kb); self.db.refresh(memory)
|
||
|
|
grant_knowledge_base(kb.id, {"user_id": self.user2.id, "permission": "use"}, self.db)
|
||
|
|
grant_memory(memory.id, {"user_id": self.user2.id, "permission": "use"}, self.db)
|
||
|
|
self.assertTrue(can_access(kb, self.user2, "read", "use"))
|
||
|
|
self.assertTrue(can_access(kb, self.user2, "use", "use"))
|
||
|
|
self.assertFalse(can_access(kb, self.user2, "manage", "use"))
|
||
|
|
self.assertEqual(len(list_admin_knowledge_bases(self.db)[0]["grants"]), 1)
|
||
|
|
self.assertEqual(len(list_admin_memories(self.db)[0]["grants"]), 1)
|
||
|
|
self.assertEqual(require_memory(self.db, memory.id, request_for(self.user2), "use").id,
|
||
|
|
memory.id)
|
||
|
|
rows = list_memories(s=self.db, request=request_for(self.user2))
|
||
|
|
self.assertEqual([row["id"] for row in rows], [memory.id])
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
unittest.main()
|