pico.chat/test.py

import asyncio
import itertools

from pico import start_server, send_json, recv_json

import websockets


class Script:

    def __init__(self):
        self.events = []

    def __add__(self, event):
        self.events.append(('send', event))
        return self

    def __sub__(self, event):
        self.events.append(('recv', event))
        return self


async def _make_client(path, timeout, script):
    state = {}
    error = False
    await asyncio.sleep(timeout)
    async with websockets.connect(path) as websocket:
        for kind, message in script.events:
            if kind == 'recv':
                A, B = message, await recv_json(websocket)
                B.pop('ts')
                state.update(message)
                if A != B:
                    error = True
                    print('-', A)
                    print('+', B)
                    print()
            elif kind == 'send':
                await send_json(websocket, **message)

        while 'users' in state:
            if min(state['users']) == state['username']:
                break
            message = await recv_json(websocket)
            message.pop('ts')
            error = True
            print('+', message)
            print()
            state.update(message)

    if error:
        print('error')
        exit()


async def _test(*clients):
    try:
        gather = asyncio.gather(
            start_server('localhost', 8642, 'TestServer'),
            *clients,
        )
        server, *results = await asyncio.wait_for(gather, timeout=2)
        server.close()
    except asyncio.TimeoutError:
        return

    return results


def test_happy_path():
    client = _make_client('ws://localhost:8642/A', 0.1, Script()
        + {'action': 'login', 'username': 'TestUser'}
        - {'kind': 'update', 'users': ['TestUser'], 'info': 'Welcome to /A', 'username': 'TestUser'}
        + {'action': 'post', 'text': 'Hello World!'}
        - {'kind': 'post', 'source': 'TestUser', 'text': 'Hello World!'}
    )
    return _test(client)


def test_name_taken():
    client1 = _make_client('ws://localhost:8642/', 0.10, Script()
        + {'action': 'login', 'username': 'A'}
        - {'kind': 'update', 'users': ['A'], 'info': 'Welcome to /', 'username': 'A'}
        - {'kind': 'update', 'users': ['A', 'B'], 'info': 'B joined'}
    )
    client2 = _make_client('ws://localhost:8642/', 0.11, Script()
        + {'action': 'login', 'username': 'A'}
        - {'kind': 'error', 'info': 'Username taken'}
    )
    client3 = _make_client('ws://localhost:8642/', 0.12, Script()
        + {'action': 'login', 'username': 'B'}
        - {'kind': 'update', 'users': ['A', 'B'], 'info': 'Welcome to /', 'username': 'B'}
        - {'kind': 'update', 'users': ['B'], 'info': 'A left'}
    )
    return _test(client1, client2, client3)


def test_interact():
    client1 = _make_client('ws://localhost:8642/', 0.10, Script()
        + {'action': 'login', 'username': 'Alice'}
        - {'kind': 'update', 'users': ['Alice'], 'info': 'Welcome to /', 'username': 'Alice'}
        - {'kind': 'update', 'users': ['Alice', 'Bob'], 'info': 'Bob joined'}
        + {'action': 'post', 'text': 'Hey Bob!'}
        - {'kind': 'post', 'source': 'Alice', 'text': 'Hey Bob!'}
        - {'kind': 'post', 'source': 'Bob', 'text': 'Howdy!'}
    )
    client2 = _make_client('ws://localhost:8642/', 0.11, Script()
        + {'action': 'login', 'username': 'Bob'}
        - {'kind': 'update', 'users': ['Alice', 'Bob'], 'info': 'Welcome to /', 'username': 'Bob'}
        - {'kind': 'post', 'source': 'Alice', 'text': 'Hey Bob!'}
        + {'action': 'post', 'text': 'Howdy!'}
        - {'kind': 'post', 'source': 'Bob', 'text': 'Howdy!'}
        + {'action': 'post', 'text': ''}
        - {'kind': 'update', 'users': ['Bob'], 'info': 'Alice left'}
    )
    return _test(client1, client2)


def test_party():
    client1 = _make_client('ws://localhost:8642/', 0.10, Script()
        + {'action': 'login', 'username': 'Norman'}
        - {'kind': 'update', 'users': ['Norman'], 'info': 'Welcome to /', 'username': 'Norman'}
        - {'kind': 'update', 'users': ['Norman', 'Ray'], 'info': 'Ray joined'}
        - {'kind': 'update', 'users': ['Emma', 'Norman', 'Ray'], 'info': 'Emma joined'}
        + {'action': 'post', 'text': 'なに？'}
        - {'kind': 'post', 'source': 'Norman', 'text': 'なに？'}
        - {'kind': 'update', 'users': ['Norman', 'Ray'], 'info': 'Emma left'}
    )
    client2 = _make_client('ws://localhost:8642/', 0.11, Script()
        + {'action': 'login', 'username': 'Ray'}
        - {'kind': 'update', 'users': ['Norman', 'Ray'], 'info': 'Welcome to /', 'username': 'Ray'}
        - {'kind': 'update', 'users': ['Emma', 'Norman', 'Ray'], 'info': 'Emma joined'}
        - {'kind': 'post', 'source': 'Norman', 'text': 'なに？'}
        - {'kind': 'update', 'users': ['Norman', 'Ray'], 'info': 'Emma left'}
        - {'kind': 'update', 'users': ['Ray'], 'info': 'Norman left'}
    )
    client3 = _make_client('ws://localhost:8642/', 0.12, Script()
        + {'action': 'login', 'username': 'Emma'}
        - {'kind': 'update', 'users': ['Emma', 'Norman', 'Ray'], 'info': 'Welcome to /', 'username': 'Emma'}
        - {'kind': 'post', 'source': 'Norman', 'text': 'なに？'}
    )
    return _test(client1, client2, client3)


def test_rooms():
    client1 = _make_client('ws://localhost:8642/A', 0.10, Script()
        + {'action': 'login', 'username': 'Dandy'}
        - {'kind': 'update', 'users': ['Dandy'], 'info': 'Welcome to /A', 'username': 'Dandy'}
        + {'action': 'post', 'text': 'Hi'}
        - {'kind': 'post', 'source': 'Dandy', 'text': 'Hi'}
    )
    client2 = _make_client('ws://localhost:8642/B', 0.10, Script()
        + {'action': 'login', 'username': 'Dandy'}
        - {'kind': 'update', 'users': ['Dandy'], 'info': 'Welcome to /B', 'username': 'Dandy'}
        + {'action': 'post', 'text': 'Howdy'}
        - {'kind': 'post', 'source': 'Dandy', 'text': 'Howdy'}
    )
    return _test(client1, client2)


if __name__ == '__main__':
    loop = asyncio.get_event_loop()
    for fn_name, fn in list(locals().items()):
        if fn_name.startswith('test_'):
            loop.run_until_complete(fn())