idealoom_load_testing.py

#!/usr/bin/env python3
import asyncio
import argparse
import pdb

from molotov import scenario, global_setup, setup_session
import simplejson as json


_SERVER = None
_HARS = []
_USER = None
_PASSWORD = None


def is_write(request):
    if request['method'] in ('GET', 'OPTIONS', 'HEAD'):
        return False
    if request['method'] != 'POST':
        return True
    if request['url'].endswith('/graphql'):
        # more for assembl for now
        data = json.loads(request['postData']['text'])
        query = data['query']
        return 'mutation' in query
    return True


def as_dict(headers, lower=False):
    def maybe_lower(h):
        return h.lower() if lower else h
    return dict(zip([maybe_lower(x['name']) for x in headers],
                    [x['value'] for x in headers]))


@scenario(weight=100)
async def from_har(session):
    global _SERVER
    har = session.har
    requests = []
    responses = []
    entries = []
    # base cleanup
    for entry in har['log']['entries']:
        request, response = entry['request'], entry['response']
        if not request['url'].startswith(_SERVER):
            continue
        if response['status'] != 200:
            continue
        resp_headers = as_dict(response['headers'], True)
        if 'cache-control' in resp_headers or 'etag' in resp_headers:
            # skip cached entries
            continue
        entries.append(entry)

    assert entries, "No request found"

    # first the writes, sequentially
    for entry in entries:
        request, response = entry['request'], entry['response']
        if not is_write(request):
            continue
        postDatum = request.get('postData', (None))
        if postDatum:
            postDatum = postDatum['text']
        resp = await session.request(
            method=request['method'],
            url=request['url'],
            params=as_dict(request['queryString']),  # assuming no multi-value
            data=postDatum,
            headers=as_dict(request['headers']),
            verify_ssl=False,
        )
        print(resp.status, request['url'])
        if resp.status not in (200, 201):
            pdb.set_trace()
        assert resp.status in (200, 201)

    # then the reads, in parallel
    for entry in entries:
        request, response = entry['request'], entry['response']
        if is_write(request):
            continue
        postDatum = request.get('postData', (None))
        if postDatum:
            postDatum = postDatum['text']
        response = session.request(
            method=request['method'],
            url=request['url'],
            params=as_dict(request['queryString']),  # assuming no multi-value
            data=postDatum,
            headers=as_dict(request['headers']),
            verify_ssl=False,
        )
        requests.append(request)
        responses.append(response)
    done, pending = await asyncio.wait(
        responses, return_when=asyncio.FIRST_EXCEPTION)
    error = False
    for task in done:
        try:
            resp = await task
            request = requests[responses.index(task._coro)]
            print(resp.status, request['url'], '\n')
            if resp.status != 200:
                print(resp.status, resp.request_info.headers)
                error = True
        except Exception as e:
            print(e)
            error = True
    assert not error


@setup_session()
async def do_login(worker_id, session):
    global _USER
    global _PASSWORD
    global _SERVER
    global _HARS
    session.har = _HARS[worker_id % len(_HARS)]
    resp1 = await session.get(_SERVER+'/about_idealoom/login')
    cookies = next(iter(session.cookie_jar._cookies.values()))
    resp = await session.post(_SERVER+'/about_idealoom/login', data={
        'referrer': 'v2',
        'identifier': _USER,
        'password': _PASSWORD
        },
        cookies=cookies)
    data = await resp.read()
    if resp.status != 200:
        print("Could not login:", resp)
        # pdb.set_trace()
        assert False, 'could not login, check credentials'
    print("cookie:", session.cookie_jar._cookies[resp.host]['assembl_session'].value)


def config(user=None, password=None, server=None, har_files=None):
    global _USER
    global _PASSWORD
    global _SERVER
    global _HARS
    if not (user and password and server and har_files):
        from configparser import ConfigParser
        parser = ConfigParser()
        with open('molotov.ini') as f:
            parser.read_file(f)
        _USER = user or parser.get('molotov', 'user')
        _PASSWORD = password or parser.get('molotov', 'password')
        _SERVER = server or parser.get('molotov', 'server')
        har_files = har_files or parser.get('molotov', 'har_files')
        assert _USER and _PASSWORD and _SERVER and har_files
    else:
        (_USER, _PASSWORD, _SERVER) = (user, password, server)
    for filename in har_files.split():
        with open(filename) as f:
            har = json.load(f)
            _HARS.append(har)


@global_setup()
def molotov_config(*args):
    # global _HARS
    config()
    # weight = 100/len(_HARS)
    # for har in _HARS:
    #     async def test(session):
    #         await from_har(session, har)
    #     scenario(weight=weight)(test)


async def main():
    import aiohttp
    async with aiohttp.ClientSession() as client:
        await do_login(0, client)
        await from_har(client)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("configuration", help="configuration file")
    parser.add_argument("-u", "--username")
    parser.add_argument("-p", "--password")
    parser.add_argument("-s", "--server")
    parser.add_argument('har', nargs='*', help='har files')
    args = parser.parse_args()
    config(args.username, args.password, args.server, args.har)
    loop = asyncio.get_event_loop()
    loop.run_until_complete(main())