test: 添加测试数据

2023-09-29 01:48:46 +08:00
parent 9bf89e5919
commit 2c3c79b21f
9 changed files with 115 additions and 5 deletions
@@ -0,0 +1,61 @@
+import csv
+import os
+
+from flask import Flask, current_app
+
+from pear_admin.extensions import db
+from pear_admin.orms import DepartmentORM, RightsORM, RoleORM, UserORM
+
+
+def dict_to_orm(d, o):
+    for k, v in d.items():
+        setattr(o, k, v or None)
+
+
+def csv_to_databases(path, orm):
+    with open(path, encoding="utf-8") as file:
+        for d in csv.DictReader(file):
+            o = orm()
+            dict_to_orm(d, o)
+            db.session.add(o)
+            db.session.flush()
+        db.session.commit()
+
+
+def register_script(app: Flask):
+    @app.cli.command()
+    def init():
+        db.drop_all()
+        db.create_all()
+
+        root = current_app.config.get("ROOT_PATH")
+
+        rights_data_path = os.path.join(root, "static", "data", "ums_rights.csv")
+        csv_to_databases(rights_data_path, RightsORM)
+
+        role_data_path = os.path.join(root, "static", "data", "ums_role.csv")
+        csv_to_databases(role_data_path, RoleORM)
+
+        with open(role_data_path, encoding="utf-8") as file:
+            for d in csv.DictReader(file):
+                role: RoleORM = RoleORM.query.get(d["id"])
+                id_list = [int(_id) for _id in d["rights_ids"].split(":")]
+                role.rights_list = RightsORM.query.filter(
+                    RightsORM.id.in_(id_list)
+                ).all()
+                db.session.commit()
+
+        department_data_path = os.path.join(
+            root, "static", "data", "ums_department.csv"
+        )
+        csv_to_databases(department_data_path, DepartmentORM)
+
+        user_data_path = os.path.join(root, "static", "data", "ums_user.csv")
+        csv_to_databases(user_data_path, UserORM)
+
+        with open(user_data_path, encoding="utf-8") as file:
+            for d in csv.DictReader(file):
+                user: UserORM = UserORM.query.get(d["id"])
+                id_list = [int(_id) for _id in d["role_ids"].split(":")]
+                user.role_list = RoleORM.query.filter(RoleORM.id.in_(id_list)).all()
+                db.session.commit()