CostEvalPlatform/test/test_login_reg.py

import unittest

from fastapi.testclient import TestClient
# from database import engine,create_db_and_tables
from sqlmodel import select, Session, SQLModel, create_engine

from config import Settings
from main import app  # 假设你的FastAPI应用实例名为app
from models import User, Tenant

settings = Settings()
engine = create_engine(settings.DATABASE_URL, echo=True)


def create_db_and_tables():
    SQLModel.metadata.create_all(engine)


# 创建一个测试客户端
client = TestClient(app)
session: Session | None = None


class TestLoginReg(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        global session
        create_db_and_tables()
        session = Session(engine)
        session.close()
        session = Session(engine)

    @classmethod
    def tearDownClass(cls):
        session.close()

    def setUp(self):
        for item in session.exec(select(User)).all():
            session.delete(item)
        session.commit()
        for item in session.exec(select(Tenant)).all():
            session.delete(item)
        session.commit()
        # 创建一个模拟的数据库会话
        # session.add()

    def tearDown(self):
        # 清理数据库
        for item in session.exec(select(User)).all():
            session.delete(item)
        session.commit()
        for item in session.exec(select(Tenant)).all():
            session.delete(item)
        session.commit()
        # session.delete()

    def test_setup_teardown(self):
        pass

    def test_login(self):
        # 创建测试用户
        test_tanant = Tenant(name="testtenant")
        session.add(test_tanant)
        session.commit()
        test_user = User(username="testuser", password="testpassword", role=1, tenant=test_tanant)
        session.add(test_user)
        session.commit()

        # 发送登录请求
        response = client.post("/api/s1/login", json={"username": "testuser", "password": "testpassword"})
        self.assertEqual(response.status_code, 200)
        self.assertEqual(response.json(), {'message': 'Login successful', 'role': 1})

    def test_register(self):
        # 发送注册请求
        response = client.post("/api/s1/register",
                               json={"name": "regtenant", "username": "reguser", "password": "regpassword"})
        self.assertEqual(response.status_code, 200)
        self.assertEqual(response.json(), {"detail": "注册成功"})

        response = client.post("/api/s1/register",
                               json={"name": "regtenant", "username": "rereguser", "password": "regpassword"})
        self.assertEqual(response.status_code, 409)
        self.assertEqual(response.json(), {"detail": "租户名已存在"})

        response = client.post("/api/s1/register",
                               json={"name": "regretenant", "username": "reguser", "password": "regpassword"})
        self.assertEqual(response.status_code, 409)
        self.assertEqual(response.json(), {"detail": "用户名已存在"})

    def get_cookie(self):
        # 创建测试用户
        test_tanant = Tenant(name="cookietenant")
        session.add(test_tanant)
        session.commit()
        test_user = User(username="cookieuser", password="cookiepassword", role=1, tenant=test_tanant)
        session.add(test_user)
        session.commit()

        # 发送登录请求以获取cookie
        login_response = client.post("/api/s1/login", json={"username": "cookieuser", "password": "cookiepassword"})
        self.assertEqual(login_response.status_code, 200)
        return login_response.cookies.get("session_token")

    def test_list_users(self):
        # 发送枚举成员请求
        cookie = self.get_cookie()
        tenant = session.exec(select(Tenant).where(Tenant.name == "cookietenant")).first()
        users = [User(username=f"testuser{i}", password="testpassword", role=2, tenant=tenant) for i in range(5)]
        session.add_all(users)
        session.commit()
        response = client.get("/api/s1/user", cookies={"session_token": cookie})
        self.assertEqual(response.status_code, 200)
        self.assertEqual(len(response.json()), 6)
        # 这里可以添加更多的断言来检查返回的用户列表

    def test_add_or_update_user(self):
        cookie = self.get_cookie()
        # 测试新增成员
        response = client.post("/api/s1/user",
                               json={"username": "newuser", "password": "newpassword", "role": "estimator"},
                               cookies={"session_token": cookie})
        self.assertEqual(response.status_code, 200)
        self.assertEqual(response.json(), {"detail": "User added successfully"})

        # 测试修改成员
        response = client.post("/api/s1/user", json={"username": "newuser", "password": "", "role": "auditor"},
                               cookies={"session_token": cookie})
        self.assertEqual(response.status_code, 200)
        self.assertEqual(response.json(), {"detail": "User updated successfully"})
        user = session.exec(select(User).where(User.username == "newuser")).first()
        self.assertEqual(user.role, 3)

    def test_delete_user(self):
        cookie = self.get_cookie()
        # 添加一个用户用于测试删除
        add_response = client.post("/api/s1/user",
                                   json={"username": "todeleteuser", "password": "deletepassword", "role": "estimator"},
                                   cookies={"session_token": cookie})
        self.assertEqual(add_response.status_code, 200)

        # 发送删除用户请求
        # 错误请求
        response = client.delete("/api/s1/user", params={"asdf": "notexistuser"}, cookies={"session_token": cookie})
        self.assertEqual(response.status_code, 422)
        # 不存在的用户
        response = client.delete("/api/s1/user", params={"username": "notexistuser"}, cookies={"session_token": cookie})
        self.assertEqual(response.status_code, 404)
        # 正确的请求
        response = client.delete("/api/s1/user", params={"username": "todeleteuser"}, cookies={"session_token": cookie})
        self.assertEqual(response.status_code, 200)
        self.assertEqual(response.json(), {"detail": "User deleted successfully"})


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