test_aio.py

# -*- coding: utf-8 -*-
import asyncio
import os
import random
import socket
import sys
# import uvloop
import threading
import time
from unittest.mock import patch

import pytest

import thriftpy2
from thriftpy2.contrib.aio.protocol import (TAsyncBinaryProtocolFactory,
                                            TAsyncCompactProtocolFactory)
from thriftpy2.contrib.aio.transport import (TAsyncBufferedTransportFactory,
                                             TAsyncFramedTransportFactory)
from thriftpy2.rpc import make_aio_client, make_aio_server
from thriftpy2.thrift import TApplicationException
from thriftpy2.transport import TTransportException

# asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())







if sys.platform == "win32":
    pytest.skip("Unix domain socket is not supported on Windows",
                allow_module_level=True)


addressbook = thriftpy2.load(os.path.join(os.path.dirname(__file__),
                                          "addressbook.thrift"))


class Dispatcher:
    def __init__(self):
        self.ab = addressbook.AddressBook()
        self.ab.people = {}

    async def ping(self):
        return True

    async def hello(self, name):
        return "hello " + name

    async def add(self, person):
        self.ab.people[person.name] = person
        return True

    async def remove(self, name):
        if not name:
            # undeclared exception
            raise ValueError('name cannot be empty')
        try:
            self.ab.people.pop(name)
            return True
        except KeyError:
            raise addressbook.PersonNotExistsError(
                "{0} not exists".format(name))

    async def get(self, name):
        try:
            return self.ab.people[name]
        except KeyError:
            raise addressbook.PersonNotExistsError(
                "{0} not exists".format(name))

    async def book(self):
        return self.ab

    async def get_phonenumbers(self, name, count):
        p = [self.ab.people[name].phones[0]] if name in self.ab.people else []
        return p * count

    async def get_phones(self, name):
        phone_numbers = self.ab.people[name].phones
        return dict((p.type, p.number) for p in phone_numbers)

    async def sleep(self, ms):
        await asyncio.sleep(ms / 1000.0)
        return True


def _create_person():
    phone1 = addressbook.PhoneNumber()
    phone1.type = addressbook.PhoneType.MOBILE
    phone1.number = '555-1212'
    phone2 = addressbook.PhoneNumber()
    phone2.type = addressbook.PhoneType.HOME
    phone2.number = '555-1234'

    # empty struct
    phone3 = addressbook.PhoneNumber()

    alice = addressbook.Person()
    alice.name = "Alice"
    alice.phones = [phone1, phone2, phone3]
    alice.created_at = int(time.time())

    return alice


class _TestAIO:
    # Base test case for all async tests
    TRANSPORT_FACTORY = NotImplemented
    PROTOCOL_FACTORY = NotImplemented

    @classmethod
    def setup_class(cls):
        cls._start_server()
        cls._start_ipv6_server()
        cls.person = _create_person()

    @classmethod
    def teardown_class(cls):
        try:
            asyncio.get_event_loop().run_until_complete(cls.server.close())
        except:  # noqa Probably already closed earlier
            pass
        del cls.server
        del cls.person

    @classmethod
    def _start_server(cls):
        cls.server = make_aio_server(
            addressbook.AddressBookService,
            Dispatcher(),
            trans_factory=cls.TRANSPORT_FACTORY,
            proto_factory=cls.PROTOCOL_FACTORY,
            loop=asyncio.new_event_loop(),
            **cls.server_kwargs(),
        )
        st = threading.Thread(target=cls.server.serve)
        st.daemon = True
        st.start()
        time.sleep(0.1)

    @classmethod
    def _start_ipv6_server(cls):
        cls.server = make_aio_server(
            addressbook.AddressBookService,
            Dispatcher(),
            trans_factory=cls.TRANSPORT_FACTORY,
            proto_factory=cls.PROTOCOL_FACTORY,
            loop=asyncio.new_event_loop(),
            socket_family=socket.AF_INET6,
            **cls.server_kwargs(),
        )
        st = threading.Thread(target=cls.server.serve)
        st.daemon = True
        st.start()
        time.sleep(0.1)

    @classmethod
    def server_kwargs(cls):
        name = cls.__name__.lower()
        return {'unix_socket': '/tmp/aio_thriftpy_test_{}.sock'.format(name)}

    @classmethod
    def client_kwargs(cls):
        return cls.server_kwargs()

    async def client(self, timeout: int = 3000000):
        return await make_aio_client(
            addressbook.AddressBookService,
            trans_factory=self.TRANSPORT_FACTORY,
            proto_factory=self.PROTOCOL_FACTORY,
            timeout=timeout,
            **self.client_kwargs(),
        )

    async def ipv6_client(self, timeout: int = 3000000):
        return await make_aio_client(
            addressbook.AddressBookService,
            trans_factory=self.TRANSPORT_FACTORY,
            proto_factory=self.PROTOCOL_FACTORY,
            timeout=timeout,
            socket_family=socket.AF_INET6,
            **self.client_kwargs(),
        )

    @pytest.mark.asyncio
    async def test_void_api(self):
        c = await self.client()
        assert await c.ping() is None
        c.close()

    @pytest.mark.asyncio
    async def test_api_ipv6(self):
        c = await self.ipv6_client()
        assert await c.ping() is None
        c.close()

    @pytest.mark.asyncio
    async def test_string_api(self):
        c = await self.client()
        assert await c.hello("world") == "hello world"
        c.close()

    @pytest.mark.asyncio
    async def test_required_argument(self):
        c = await self.client()
        assert await c.hello("") == "hello "

        with pytest.raises(TApplicationException):
            await c.hello()
        c.close()

    @pytest.mark.asyncio
    async def test_huge_res(self):
        c = await self.client()
        big_str = "world" * 100000
        assert await c.hello(big_str) == "hello " + big_str
        c.close()

    @pytest.mark.asyncio
    async def test_tstruct_req(self):
        c = await self.client()
        assert await c.add(self.person) is True
        c.close()

    @pytest.mark.asyncio
    async def test_tstruct_res(self):
        c = await self.client()
        assert self.person == await c.get("Alice")
        c.close()

    @pytest.mark.asyncio
    async def test_complex_tstruct(self):
        c = await self.client()
        assert len(await c.get_phonenumbers("Alice", 0)) == 0
        assert len(await c.get_phonenumbers("Alice", 1000)) == 1000
        c.close()

    @pytest.mark.asyncio
    async def test_exception(self):
        c = await self.client()
        with pytest.raises(addressbook.PersonNotExistsError):
            await c.remove("Bob")
        c.close()

    @pytest.mark.asyncio
    async def test_undeclared_exception(self):
        c = await self.client()
        with pytest.raises(TTransportException):
            await c.remove('')
        c.close()

    @pytest.mark.asyncio
    async def test_client_socket_timeout(self):
        c = await self.client(timeout=500)
        with pytest.raises(asyncio.TimeoutError):
            await c.sleep(1000)
        c.close()


class SSLServerMixin:

    @classmethod
    def setup_class(cls):
        cls.port = random.randint(55000, 56000)
        super().setup_class()

    @classmethod
    def server_kwargs(cls):
        return {
            'host': 'localhost',
            'port': cls.port,
            'certfile': "ssl/server.pem",
            'keyfile': "ssl/server.key",
        }

    @classmethod
    def client_kwargs(cls):
        kw = cls.server_kwargs()
        kw['cafile'] = "ssl/CA.pem"
        return kw

    async def client_with_url(self, timeout: int = 3000):
        kw = self.client_kwargs()
        kw['url'] = "thrift://{}:{}".format(kw.pop('host'), kw.pop('port'))
        return await make_aio_client(
            addressbook.AddressBookService,
            trans_factory=self.TRANSPORT_FACTORY,
            proto_factory=self.PROTOCOL_FACTORY,
            timeout=timeout,
            **kw,
        )

    @pytest.mark.asyncio
    async def test_clients(self):
        c1 = await self.client()
        c2 = await self.client_with_url()
        assert await c1.hello("world") == await c2.hello("world")
        c1.close()
        c2.close()


class TestAIOBufferedBinary(_TestAIO):
    TRANSPORT_FACTORY = TAsyncBufferedTransportFactory()
    PROTOCOL_FACTORY = TAsyncBinaryProtocolFactory()


class TestAIOBufferedCompact(_TestAIO):
    TRANSPORT_FACTORY = TAsyncBufferedTransportFactory()
    PROTOCOL_FACTORY = TAsyncCompactProtocolFactory()


class TestAIOFramedBinary(_TestAIO):
    TRANSPORT_FACTORY = TAsyncFramedTransportFactory()
    PROTOCOL_FACTORY = TAsyncBinaryProtocolFactory()


class TestAIOFramedCompact(_TestAIO):
    TRANSPORT_FACTORY = TAsyncFramedTransportFactory()
    PROTOCOL_FACTORY = TAsyncCompactProtocolFactory()


class TestAIOBufferedBinarySSL(SSLServerMixin, TestAIOBufferedBinary):
    pass


class TestAIOBufferedCompactSSL(SSLServerMixin, TestAIOBufferedCompact):
    pass


class TestAIOFramedBinarySSL(SSLServerMixin, TestAIOFramedBinary):
    pass


class TestAIOFramedCompactSSL(SSLServerMixin, TestAIOFramedCompact):
    pass


@pytest.mark.asyncio
async def test_client_connect_timeout():
    with pytest.raises(TTransportException):
        c = await make_aio_client(
            addressbook.AddressBookService,
            unix_socket='/tmp/test.sock',
            connect_timeout=1000
        )
        await c.hello('test')


class TestDeprecatedTimeoutKwarg:
    """
    Replace TAsyncSocket with a Mock object that raises a RuntimeError
    when called. This allows us to check that timeout vs. socket_timeout
    arguments are properly handled without actually creating the client.

    This class should be removed when the socket_timeout argument is removed.
    """
    def setup_method(self):
        # Create and apply a fresh patch for each test.
        self.async_sock = patch(
            'thriftpy2.contrib.aio.rpc.TAsyncSocket',
            side_effect=RuntimeError,
        ).__enter__()

    def teardown_method(self):
        self.async_sock.__exit__()  # Clean up patch

    @pytest.mark.asyncio
    async def test_no_timeout_given(self):
        await self._make_client()
        assert self._given_timeout() == 3000  # Default value

    @pytest.mark.asyncio
    async def test_timeout_given(self):
        await self._make_client(timeout=1234)
        assert self._given_timeout() == 1234

    @pytest.mark.asyncio
    async def test_socket_timeout_given(self):
        await self._make_client(warning=DeprecationWarning, socket_timeout=555)
        assert self._given_timeout() == 555

    @staticmethod
    async def _make_client(warning=None, **kwargs):
        """
        Helper method to create the client and check that the proper warning
        is emitted (if any) and that the patch is properly applied by
        consuming the RuntimeError.
        """
        if warning:
            with pytest.warns(warning),\
                    pytest.raises(RuntimeError):  # Consume error
                await make_aio_client(addressbook.AddressBookService, **kwargs)
        else:
            with pytest.raises(RuntimeError):  # Consume error
                await make_aio_client(addressbook.AddressBookService, **kwargs)

    def _given_timeout(self):
        """Get the timeout provided to TAsyncSocket."""
        try:
            self.async_sock.assert_called_once()
        except AttributeError:  # Python 3.5
            assert self.async_sock.call_count == 1
        _args, kwargs = self.async_sock.call_args
        return kwargs['socket_timeout']