性能测试 - Locust TCP socket client（转载）

一、背景

Locust 是性能测试工具，但是默认只支持http协议，就是默认只有http的client，需要其他协议的测试必须自己扩展对于的client，比如下面的TCP client。

二、二次开发（类重写）

1. Tcp socket client 类

import time
import random
# from socket import socket, AF_INET, SOCK_STREAM
import socket
from locust import Locust, TaskSet, events, task
 
# author: Max.Bai
# date: 2017-04
 
class TcpSocketClient(socket.socket):
    # locust tcp client
    # author: Max.Bai@2017
    def __init__(self, af_inet, socket_type):
        super(TcpSocketClient, self).__init__(af_inet, socket_type)
 
    def connect(self, addr):
        start_time = time.time()
        try:
            super(TcpSocketClient, self).connect(addr)
        except Exception as e:
            total_time = int((time.time() - start_time) * 1000)
            events.request_failure.fire(request_type="tcpsocket", name="connect", response_time=total_time, exception=e)
        else:
            total_time = int((time.time() - start_time) * 1000)
            events.request_success.fire(request_type="tcpsocket", name="connect", response_time=total_time, response_length=0)
        
    def send(self, msg):
        start_time = time.time()
        try:
            super(TcpSocketClient, self).send(msg)
        except Exception as e:
            total_time = int((time.time() - start_time) * 1000)
            events.request_failure.fire(request_type="tcpsocket", name="send", response_time=total_time, exception=e)
        else:
            total_time = int((time.time() - start_time) * 1000)
            events.request_success.fire(request_type="tcpsocket", name="send", response_time=total_time, response_length=0)
 
    def recv(self, bufsize):
        recv_data = ''
        start_time = time.time()
        try:
            recv_data = super(TcpSocketClient, self).recv(bufsize)
        except Exception as e:
            total_time = int((time.time() - start_time) * 1000)
            events.request_failure.fire(request_type="tcpsocket", name="recv", response_time=total_time, exception=e)
        else:
            total_time = int((time.time() - start_time) * 1000)
            events.request_success.fire(request_type="tcpsocket", name="recv", response_time=total_time, response_length=0)
        return recv_data

2. Locust 类

class TcpSocketLocust(Locust):
    """
    This is the abstract Locust class which should be subclassed. It provides an TCP socket client
    that can be used to make TCP socket requests that will be tracked in Locust's statistics.
    author: Max.bai@2017
    """
    def __init__(self, *args, **kwargs):
        super(TcpSocketLocust, self).__init__(*args, **kwargs)
        self.client = TcpSocketClient(socket.AF_INET, socket.SOCK_STREAM)
        ADDR = (self.host, self.port)
        self.client.connect(ADDR)

三、Locust-Tcp应用

1、Locust运行应用

class TcpTestUser(TcpSocketLocust):
    
    host = "127.0.0.1"
    port = 12345
    min_wait = 100
    max_wait = 1000
 
    class task_set(TaskSet):        
        @task
        def send_data(self):
            self.client.send(random_str())
            data = self.client.recv(2048).decode()
            print(data)
 
 
if __name__ == "__main__":
    user = TcpTestUser()
    user.run()

2、Tcp server 方便测试使用

#!/usr/bin/env python3
#-*- coding:utf-8 -*-
 
import time
from socket import socket, AF_INET, SOCK_STREAM
import threading
 
buffsize = 2048
 
def tcplink(sock, addr):
    # print 'Accept new connection from %s:%s...' % addr
    sock.send('Welcome!')
    while True:
        try:
            # data = sock.recv(buffsize)
            data = sock.recv(buffsize).decode()
            time.sleep(1)
            if data == 'exit' or not data:
                break
            sock.send('Hello, %s!' % data)
        except Exception as e:
            print(str(e))
            break
    sock.close()
    print 'Connection from %s:%s closed.' % addr
 
def main():
    host = ''
    port = 12345
 
    ADDR = (host,port)
 
    tctime = socket(AF_INET,SOCK_STREAM)
    tctime.bind(ADDR)
    tctime.listen(3)
 
    print('Wait for connection ...')
    while True:
        sock,addr = tctime.accept()
        print 'Accept new connection from %s:%s...' % addr
 
        t = threading.Thread(target=tcplink, args=(sock, addr))
        t.start()
 
if __name__=="__main__":
    main()