# locustfile.py
import random
import string

from locust import HttpUser, TaskSet, task, between

class UserBehavior(TaskSet):

    def on_start(self):
        """在测试开始时执行，进行登录并存储 token 和 username"""
        # 假设有一个登录端点，可以获取 token 和 username
        response = self.client.post("http://127.0.0.1:8000/teacher/jwt/token", data={
            "username": "starhun",
            "password": "123456"
        })
        if response.status_code == 200:
            self.token = response.json()["access_token"]
            self.username = response.json().get("username")
        else:
            self.token = None
            self.username = None

    @task(1)
    def get_classes(self):
        """测试获取班级列表"""
        if self.token:
            with self.client.get("http://127.0.0.1:8000/teacher/pick_student/getclasses", headers={"Authorization": f"Bearer {self.token}"}) as response:
                if response.status_code != 200:
                    response.failure(f"Failed to get classes: {response.text}")

    @task(2)
    def create_class(self):
        characters = string.ascii_letters + string.digits
        # 使用 random.choice() 从字符集合中随机选择字符
        random_string = ''.join(random.choice(characters) for i in range(10))
        """测试创建班级"""
        if self.token:
            with self.client.post("http://127.0.0.1:8000/teacher/pick_student/create_class", params={"class_name": random_string,"class_time":"周三 5-6节"}, headers={"Authorization": f"Bearer {self.token}"}, catch_response=True) as response:
                if response.status_code != 200:
                    response.failure(f"Failed to create class: {response.text}")

    @task(3)
    def random_pick(self):
        """测试随机点名"""
        if self.token:
            params = {"class_name": "aa"}
            with self.client.get("http://127.0.0.1:8000/teacher/pick_student/random_pick/", params=params, headers={"Authorization": f"Bearer {self.token}"}, catch_response=True) as response:
                if response.status_code != 200:
                    response.failure(f"Failed to random pick: {response.text}")

    @task(4)
    def upload_students(self):
        """测试上传学生信息"""
        if self.token:
            files = {
                "file": ("stduents.xlsx", open("D:\python\code\pick_student\stduents.xlsx", "rb"), "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet")
            }
            data = {"class_name": "班级1"}
            with self.client.post("http://127.0.0.1:8000/teacher/pick_student/upload_students/", files=files, params=data, headers={"Authorization": f"Bearer {self.token}"}, catch_response=True) as response:
                if response.status_code != 200:
                    response.failure(f"Failed to upload students: {response.text}")

class WebsiteUser(HttpUser):
    tasks = [UserBehavior]
    wait_time = between(1, 5)