python实现websocket测试脚本

在此记录一下学习过程！
参考了:https://blog.csdn.net/ennismar/article/details/78190887
和https://segmentfault.com/q/1010000009284816?_ea=1883181
想实现多线程并发的socket的测试脚本！

#coding:utf-8
from ws4py.client.threadedclient import WebSocketClient
from threadpool import ThreadPool,makeRequests
import time,threading,requests,json
from framework import logger2


logger = logger2.Logger(logger='test').getlog()

# 开房人ID从201开始为第一个！
create_room_user_num = 201

#房间数量
num = 10
#开房人数量
people_num = 3
#房间人数
np = 6
"""
        if np == 6:
            people_num = 0
        elif np == 9:
            people_num = 1
        elif np == 10:
            people_num = 2
        elif np == 12:
            people_num = 3
        elif np == 13:
            people_num = 4
"""
#线程数量=房间数量*房间人数*开房人数量
thread_num = num*np*people_num

#开房者ID

userid_man = 'test281'
home = 'XXX'
ws_home = 'XXX'

class Login():
    def __init__(self,user_id):
        self.user_id = user_id

    def login(self):
        login_params = {
            'userId': self.user_id,
            'password': ''
        }

        login_url = home + 'login2'
        login_re = requests.post(login_url, params=login_params).json()
        token = login_re['data']['token']
        user_id = login_re['data']['userId']

        login = {
            'token':token,
            'user_id':user_id
        }

        return login


class OpenRoom():
    def __init__(self,people_num,user_id,room_num=1):
        self.user_id = user_id
        self.room_num = room_num
        self.people_num = people_num

    def open_room(self):
        global people_num

        if self.people_num == 6:
            people_num = 0
        elif self.people_num == 9:
            people_num = 1
        elif self.people_num == 10:
            people_num = 2
        elif self.people_num == 12:
            people_num = 3
        elif self.people_num == 13:
            people_num = 4

        token = Login(self.user_id).login()['token']

        open_room_params_createRoomV0 = json.dumps(
            {
            "creatorId":self.user_id,
            "gameType": "NIUNIU",
            "dipai": 1,
            "playingType": 0,
            "endPoints": 4,
            "numberGames": 2,
            "peopleNum": people_num,
            "rule": 1,
            "bankerMultiple": 1,
            "setoutTime": 5,
            "campaignTime": 5,
            "betTime": 5,
            "showdownTime": 5,
            "automaticMode": 1,
            "tableNumber": 0,
            "sihua": 0,
            "shunzi": 0,
            "tonghua": 0,
            "wuhua": 1,
            "hulu": 0,
            "wuxiao": 1,
            "zhadan": 1,
            "tonghuashun": 0,
            "chip_1": 0,
            "chip_2": 0,
            "chip_3": 0,
            "chip_4": 0,
            "chip_5": 1,
            "chip_6": 1,
            "chip_8": 1,
            "chip_10": 1
            }
        )

        open_room_params = {
            'token':token,
            'createRoomVO':open_room_params_createRoomV0
        }

        open_room_url = home + 'createRoom'

        room_id = []

        for _ in range(self.room_num):
            open_room_re = requests.post(open_room_url,params=open_room_params).json()

            if open_room_re['status'] == 1074:
                logger.error(open_room_re)
                logger.error('=====不能创建房间，退出程序！=====')
                exit()
            else:
                room = open_room_re['data']['roomId']
                room_id.append(room)
        logger.info('创建的房间列表：%s' % room_id)
        return room_id

class AddGame():
    def __init__(self,user_id,room):
        self.user_id = 'test3' + str(user_id)
        self.room = room

    def add_game(self):
        login = Login(self.user_id).login()
        user_id = login['user_id']
        token = login['token']

        add_game_param = json.dumps(
            {
            "playerId":user_id,
            "tableId": 1,
            "roomId":self.room
            }
        )

        add_game_params = {
            'token':token,
            'addGameParam':add_game_param
        }

        url = home + 'addGame'
        re = requests.post(url,params=add_game_params).json()
        logger.debug('add_game:%s'% re)

        add_game = {
            'user_id':user_id,
            'token':token
        }
        return add_game


class DummyClient(WebSocketClient):

    def __init__(self, url, user_id,token,room,num,protocols=None, extensions=None, heartbeat_freq=None,
                 ssl_options=None, headers=None, exclude_headers=None):
        WebSocketClient.__init__(self, url, protocols=None, extensions=None, heartbeat_freq=None,
                 ssl_options=None, headers=None, exclude_headers=None)
        self.user_id = user_id
        self.token = token
        self.room = room
        self.num = num

    def opened(self):

        def send_thread():
            """
            101是确定第一个包，确认链接
            105是坐下
            114是准备
            :return:
            """
            send_keep = "ping"

            join_info = json.dumps({
                "userId": self.user_id,
                "token": self.token,
                "roomId": self.room,
                "tableId": "1"
            })

            send_join = json.dumps({
                "cammand": 101,
                "data":join_info
            })

            send_sit = json.dumps({
                "cammand": 105
            })

            send_ready = json.dumps({
                "cammand": 114
            })

            self.send(send_join)
            self.send(send_keep)
            time.sleep(3)
            self.send(send_sit)

            #等待10秒后，发出114（准备），已便等待其他玩家准备
            # for _ in range(5):
            #     self.send(send_keep)
            #     time.sleep(4)

            # self.send(send_ready)
            while True:
                time.sleep(5)
                self.send(send_keep)

        t = threading.Thread(target=send_thread)
        time.sleep(1)
        t.start()
        time.sleep(3)


    def closed(self, code, reason=None):
        logger.error("Closed down:%s,%s" % (code, reason))

    def received_message(self, m):
        logger.info(str(self.user_id) +' = '+ str(m))

def on_start(num):
    time.sleep(10)

    room_num = room_id_list[num//np]

    b = AddGame(num,room_num).add_game()

    ws = DummyClient(ws_home + 'web_sk',
                     protocols=['chat'],
                     user_id=b['user_id'],
                     token=b['token'],
                     room=room_num,
                     num=num
                     )

    try:
        ws.connect()
    except:
        logger.error('connect！=====================')
    try:
        ws.run_forever()
    except:
        logger.error("run_forever!===================")


if __name__ == "__main__":
    pool = ThreadPool(thread_num)

    test = list()
    for ij in range(thread_num):
        test.append(ij)

    room_id_list = []


    for _ in range(people_num):
        create_room_user_id = 'test' + str(create_room_user_num)
        room_id = OpenRoom(np,create_room_user_id,num).open_room()
        room_id_list.extend(room_id)
        create_room_user_num += 1

    re = makeRequests(on_start,test)
    [pool.putRequest(req) for req in re]

    pool.wait()

来源：CSDN

作者：choy_cc

链接：https://blog.csdn.net/choy_cc/article/details/82462493