|
| 1 | +import json |
| 2 | +import os |
| 3 | +from random import choice |
| 4 | + |
| 5 | +import django |
| 6 | +import requests |
| 7 | +from django.db.models import Q |
| 8 | +from locust import between, task |
| 9 | +from locust_plugins.users.socketio import SocketIOUser |
| 10 | +from websocket import WebSocketConnectionClosedException |
| 11 | + |
| 12 | +# Django 설정 모듈 지정 |
| 13 | +os.environ["DJANGO_SETTINGS_MODULE"] = "config.settings.settings" |
| 14 | + |
| 15 | +# Django 초기화 |
| 16 | +django.setup() |
| 17 | + |
| 18 | +from apps.chat.models import Chatroom |
| 19 | + |
| 20 | + |
| 21 | +def get_test_chatroom_info(): # type: ignore |
| 22 | + # Q 객체를 사용하여 OR 조건으로 필터링 |
| 23 | + |
| 24 | + filtered_chatrooms = Chatroom.objects.filter( |
| 25 | + Q(borrower__nickname__icontains="test") | Q(lender__nickname__icontains="test") |
| 26 | + ) |
| 27 | + # 필터링된 채팅방의 ID, 대여자, 판매자 목록을 list[dict]로 반환 |
| 28 | + return filtered_chatrooms.values("id", "borrower", "lender") |
| 29 | + |
| 30 | + |
| 31 | +from apps.user.models import Account |
| 32 | + |
| 33 | + |
| 34 | +def get_account_model(user_id: int) -> Account: |
| 35 | + |
| 36 | + return Account.objects.get(id=user_id) |
| 37 | + |
| 38 | + |
| 39 | +# 채팅방 데이터 리스트 |
| 40 | +chatrooms = get_test_chatroom_info() # type: ignore |
| 41 | + |
| 42 | + |
| 43 | +class WebSocketUser(SocketIOUser): |
| 44 | + wait_time = between(1, 5) # type: ignore |
| 45 | + |
| 46 | + def on_start(self) -> None: |
| 47 | + # 랜덤한 채팅방과 유저를 선택 |
| 48 | + chatroom_info = choice(chatrooms) |
| 49 | + self.chatroom_id = chatroom_info["id"] |
| 50 | + self.borrower = get_account_model(chatroom_info["borrower"]) |
| 51 | + self.lender = get_account_model(chatroom_info["lender"]) |
| 52 | + self.sender = choice([self.borrower, self.lender]) |
| 53 | + # 선택된 유저로 로그인 요청 보내고 csrf, access 토큰 가져오기 |
| 54 | + login_req = requests.post( |
| 55 | + f"https://{os.environ.get("BACKEND_HOST")}/api/users/login/", |
| 56 | + data={"email": self.sender.email, "password": "password123@"} |
| 57 | + ) |
| 58 | + csrf = login_req.cookies.get("csrftoken") |
| 59 | + session_cookie = login_req.cookies.get("sessionid") |
| 60 | + access = login_req.cookies.get("adfdfd") |
| 61 | + print(f"csrf: {csrf}\n, session: {session_cookie}\n, access: {access}") |
| 62 | + |
| 63 | + print(f"chatroom_id: {self.chatroom_id}, borrower: {self.borrower.nickname}, lender: {self.lender.nickname}") |
| 64 | + |
| 65 | + if csrf and access and session_cookie: |
| 66 | + self.connect( |
| 67 | + f"ws://{os.environ.get("BACKEND_HOST")}/ws/chat/{self.chatroom_id}/", |
| 68 | + header=[f"X-CSRFToken: {csrf}", f"Authorization: Bearer {access}"], |
| 69 | + cookie="sessionid=" + session_cookie, |
| 70 | + ) |
| 71 | + print( |
| 72 | + f"Connected to chatroom_id: {self.chatroom_id} as borrower: {self.borrower.nickname}, lender: {self.lender.nickname}" |
| 73 | + ) |
| 74 | + else: |
| 75 | + raise ValueError("로그인 요청 실패") |
| 76 | + |
| 77 | + @task |
| 78 | + def send_message(self) -> None: |
| 79 | + if self.sender == self.borrower: |
| 80 | + message = { |
| 81 | + "text": f"Hello! i want borrow your clothes!", |
| 82 | + "sender": self.sender.nickname, |
| 83 | + } |
| 84 | + self.ws.send(json.dumps(message)) |
| 85 | + if self.sender == self.lender: |
| 86 | + message = { |
| 87 | + "text": f"Of course! When are you going to borrow and return it?", |
| 88 | + "sender": self.sender.nickname, |
| 89 | + } |
| 90 | + self.ws.send(json.dumps(message)) |
| 91 | + |
| 92 | + def on_message(self, message: bytes) -> None: |
| 93 | + try: |
| 94 | + if message: |
| 95 | + response = json.loads(message) |
| 96 | + print(f"Received: {response}") |
| 97 | + except json.JSONDecodeError: |
| 98 | + print(f"Received message is not valid JSON") |
| 99 | + |
| 100 | + def on_stop(self) -> None: |
| 101 | + self.ws.close() |
| 102 | + |
| 103 | + def receive_loop(self) -> None: |
| 104 | + while True: |
| 105 | + try: |
| 106 | + message = self.ws.recv() |
| 107 | + import logging |
| 108 | + |
| 109 | + logging.debug(f"WSR: {message}") |
| 110 | + self.on_message(message) |
| 111 | + except WebSocketConnectionClosedException as e: |
| 112 | + print(f"소켓 연결 종료됨 \nMessage :{e}") |
| 113 | + break |
0 commit comments