Test mass connection to the openslides projector

openslides_ws_test.py

#!/usr/bin/env python

import asyncio
import functools
from datetime import datetime, timedelta

import websockets

WEBSOCKET_URI = 'ws://localhost:8000/ws/projector/1/'
CONNECTION_COUNT = 500


async def mass_connect():
    object_list = list()
    task_list = list()
    connection_list = list()

    def recv_callback(obj, future):
        """
        Get called when a websocket connection recv data.

        Saves the data and the current time
        """
        obj['recv_time'] = datetime.now()

    # open connections
    for i in range(CONNECTION_COUNT):
        # Connect to the server. connection is a WebSocketClientProtocol. Waits
        # until the connection is established.
        connection = await websockets.connect(WEBSOCKET_URI)
        connection_list.append(connection)

        # Create a dictonary with the connection. Later, some meta data will be
        # saved into this dict
        recv_object = {
            'connection': connection,
            'start_time': datetime.now(),
            'recv_time': None}

        # Create a task that is done, when the connection receives data
        task = asyncio.ensure_future(connection.recv())

        # Call recv_callback with the argument recv_object when the connection
        # got data.
        task.add_done_callback(functools.partial(recv_callback, recv_object))

        # Save the recv_object and the task for later useage
        object_list.append(recv_object)
        task_list.append(task)

    # Wait for all connections to receive the changed data. Pending is a list
    # of tasks that did not receive any data until timeout
    __, pending = await asyncio.wait(task_list, timeout=10)

    max_time = timedelta.min
    min_time = timedelta.max
    close_tasks = set()
    for recv_object in object_list:
        # Get the max_time and the min_time
        if recv_object['recv_time'] is not None:
            elapsed_time = recv_object['recv_time'] - recv_object['start_time']
            max_time = max(max_time, elapsed_time)
            min_time = min(min_time, elapsed_time)

        # Close the connection
        close_tasks.add(asyncio.ensure_future(recv_object['connection'].close()))

    # Print some status informations
    print("Max: {}\nMin: {}\nFail: {}\n".format(
        max_time, min_time, len(pending)))

    # Wait until all connections are closed
    await asyncio.wait(close_tasks)


def exception_handler(loop, context):
    if isinstance(context['exception'], websockets.exceptions.ConnectionClosed):
        pass
    else:
        loop.default_exception_handler(context)

loop = asyncio.get_event_loop()
loop.set_exception_handler(exception_handler)
loop.run_until_complete(mass_connect())