zkxw/backend/tests/test_user_list_relations_perf.py

"""用户列表关联补全：须批量查询，避免按用户 N+1。"""
import unittest
from datetime import datetime

from sqlalchemy import create_engine, event
from sqlalchemy.orm import sessionmaker


class TestEnrichUserRelationsPerf(unittest.TestCase):
    def setUp(self):
        from app.models.models import (
            SysPost,
            SysPostRole,
            SysRole,
            SysUser,
            SysUserPost,
            SysUserRole,
        )

        self.engine = create_engine("sqlite:///:memory:")
        for table in (
            SysUser.__table__,
            SysRole.__table__,
            SysUserRole.__table__,
            SysPost.__table__,
            SysUserPost.__table__,
            SysPostRole.__table__,
        ):
            table.create(self.engine)
        Session = sessionmaker(bind=self.engine, autoflush=False, autocommit=False)
        self.db = Session()
        self.now = datetime.now()
        self.query_count = 0

        @event.listens_for(self.engine, "before_cursor_execute")
        def _count(_conn, _cursor, _statement, _parameters, _context, _executemany):
            self.query_count += 1

        self._listener = _count
        self.SysUser = SysUser
        self.SysRole = SysRole
        self.SysUserRole = SysUserRole
        self.SysPost = SysPost
        self.SysUserPost = SysUserPost
        self.SysPostRole = SysPostRole

    def tearDown(self):
        event.remove(self.engine, "before_cursor_execute", self._listener)
        self.db.close()
        self.engine.dispose()

    def _seed(self, n: int = 20):
        self.db.add(
            self.SysRole(
                id="role-direct",
                role_name="运营管理员",
                status="1",
                tenant_no="zkxw",
                create_dept_id="",
                create_time=self.now,
                update_time=self.now,
                del_flag=0,
            )
        )
        self.db.add(
            self.SysRole(
                id="role-post",
                role_name="职位角色",
                status="1",
                tenant_no="zkxw",
                create_dept_id="",
                create_time=self.now,
                update_time=self.now,
                del_flag=0,
            )
        )
        self.db.add(
            self.SysPost(
                id="post-1",
                post_name="人事",
                post_code="hr",
                status="1",
                tenant_no="zkxw",
                create_dept_id="",
                create_time=self.now,
                update_time=self.now,
                del_flag=0,
            )
        )
        self.db.add(
            self.SysPostRole(
                id="pr-1",
                post_id="post-1",
                role_id="role-post",
                tenant_no="zkxw",
                create_dept_id="",
                create_time=self.now,
                update_time=self.now,
                del_flag=0,
            )
        )
        for i in range(n):
            uid = f"u{i}"
            self.db.add(
                self.SysUser(
                    id=uid,
                    user_no=f"user{i}",
                    nick_name=f"用户{i}",
                    password="x",
                    status="1",
                    tenant_no="zkxw",
                    create_dept_id="",
                    create_time=self.now,
                    update_time=self.now,
                    del_flag=0,
                )
            )
            self.db.add(
                self.SysUserRole(
                    id=f"ur{i}",
                    user_id=uid,
                    role_id="role-direct",
                    tenant_no="zkxw",
                    create_dept_id="",
                    create_time=self.now,
                    update_time=self.now,
                    del_flag=0,
                )
            )
            self.db.add(
                self.SysUserPost(
                    id=f"up{i}",
                    user_id=uid,
                    post_id="post-1",
                    tenant_no="zkxw",
                    create_dept_id="",
                    create_time=self.now,
                    update_time=self.now,
                    del_flag=0,
                )
            )
        self.db.commit()

    def test_enrich_query_count_bounded_and_fields_correct(self):
        from app.crud import user as user_crud

        n = 20
        self._seed(n)
        items = [{"id": f"u{i}"} for i in range(n)]
        self.query_count = 0
        out = user_crud.enrich_user_dicts_with_relations(self.db, items)

        self.assertIs(out, items)
        self.assertEqual(len(items), n)
        first = items[0]
        self.assertEqual(first["roleIds"], ["role-direct"])
        self.assertEqual(first["postIds"], ["post-1"])
        self.assertEqual(first["postInheritedRoleIds"], ["role-post"])
        self.assertEqual(first["effectiveRoleIds"], ["role-direct", "role-post"])
        self.assertEqual(first["effectiveRoleNames"], ["运营管理员", "职位角色"])
        # 批量查询：角色/职位关联 + 启用角色名，不应随用户数线性增长
        self.assertLessEqual(self.query_count, 8)
        self.assertLess(self.query_count, n)


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