async marionette calls

marionette.py

# see https://dxr.mozilla.org/mozilla-central/source/testing/marionette/driver.js?from=testing%2Fmarionette%2Fdriver.js#2957
import functools
import json
from asyncio import open_connection
import asyncio
from molotov import *


class Marionette(object):
    def __init__(self, host='localhost', port=2828, loop=None):
        self.host = host
        self.port = port
        self.loop = loop
        self._r = self._w = None

    async def __aenter__(self):
        await self.open()
        return self

    async def open(self):
        self._r, self._w = await open_connection(self.host, self.port,
                                                 loop=self.loop)
        await self.read()
        await self('newSession')

    async def close(self):
        await self('deleteSession')
        self._w.close()
        self._r = self._w = None

    async def __aexit__(self, exc_type, exc, tb):
        await self.close()

    async def read(self):
        reader = self._r
        size = (await reader.readuntil(b':')).decode()
        size = int(size[:-1])
        data = (await reader.readexactly(size)).decode()
        return json.loads(data)

    async def send(self, command, **options):
        mid = 1
        message = [0, mid, command, options]
        message = json.dumps(message)
        message = '%d:%s' % (len(message), message)
        self._w.write(message.encode())
        data = await self.read()
        return data

    __call__ = send



@session_setup()
async def _setup(session):
    session.browser = Marionette(loop=session.loop)
    await session.browser.open()


@session_teardown()
async def _teardown(session):
    await session.browser.close()


@scenario(1)
async def test(session):
    await session.browser('getSessionCapabilities')
    await session.browser('get', url='http://ziade.org')