Перейти к содержимому

Назад

test_utils.py

Lib/asyncio/test_utils.pyCPython 3.6
python
"""Утилиты, общие для тестов."""

import collections
import contextlib
import io
import logging
import os
import re
import socket
import socketserver
import sys
import tempfile
import threading
import time
import unittest
import weakref

from unittest import mock

from http.server import HTTPServer
from wsgiref.simple_server import WSGIRequestHandler, WSGIServer

try:
    import ssl
except ImportError:  # pragma: no cover
    ssl = None

from . import base_events
from . import compat
from . import events
from . import futures
from . import selectors
from . import tasks
from .coroutines import coroutine
from .log import logger
from test import support


if sys.platform == 'win32':  # pragma: no cover
    from .windows_utils import socketpair
else:
    from socket import socketpair  # pragma: no cover


def data_file(filename):
    if hasattr(support, 'TEST_HOME_DIR'):
        fullname = os.path.join(support.TEST_HOME_DIR, filename)
        if os.path.isfile(fullname):
            return fullname
    fullname = os.path.join(os.path.dirname(os.__file__), 'test', filename)
    if os.path.isfile(fullname):
        return fullname
    raise FileNotFoundError(filename)


ONLYCERT = data_file('ssl_cert.pem')
ONLYKEY = data_file('ssl_key.pem')


def dummy_ssl_context():
    if ssl is None:
        return None
    else:
        return ssl.SSLContext(ssl.PROTOCOL_SSLv23)


def run_briefly(loop):
    @coroutine
    def once():
        pass
    gen = once()
    t = loop.create_task(gen)
    # Не логировать предупреждение, если задача не завершена после run_until_complete().
    # Это происходит, если цикл остановлен или задача вызывает BaseException.
    t._log_destroy_pending = False
    try:
        loop.run_until_complete(t)
    finally:
        gen.close()


def run_until(loop, pred, timeout=30):
    deadline = time.time() + timeout
    while not pred():
        if timeout is not None:
            timeout = deadline - time.time()
            if timeout <= 0:
                raise futures.TimeoutError()
        loop.run_until_complete(tasks.sleep(0.001, loop=loop))


def run_once(loop):
    """Устаревший API для однократного запуска через цикл событий.

    Это рекомендуемый шаблон для тестового кода.  Он опрашивает
    селектор один раз и запускает все колбэки, запланированные в ответ на события
    ввода-вывода.
    """
    loop.call_soon(loop.stop)
    loop.run_forever()


class SilentWSGIRequestHandler(WSGIRequestHandler):

    def get_stderr(self):
        return io.StringIO()

    def log_message(self, format, *args):
        pass


class SilentWSGIServer(WSGIServer):

    request_timeout = 2

    def get_request(self):
        request, client_addr = super().get_request()
        request.settimeout(self.request_timeout)
        return request, client_addr

    def handle_error(self, request, client_address):
        pass


class SSLWSGIServerMixin:

    def finish_request(self, request, client_address):
        # Относительное расположение нашего тестового каталога (который
        # содержит файлы ключей и сертификатов SSL) отличается
        # между stdlib и автономным asyncio.
        # Предпочитать наш собственный, если удается его найти.
        keyfile = ONLYKEY
        certfile = ONLYCERT
        context = ssl.SSLContext()
        context.load_cert_chain(certfile, keyfile)

        ssock = context.wrap_socket(request, server_side=True)
        try:
            self.RequestHandlerClass(ssock, client_address, self)
            ssock.close()
        except OSError:
            # возможно, сокет был закрыт одноранговым узлом
            pass


class SSLWSGIServer(SSLWSGIServerMixin, SilentWSGIServer):
    pass


def _run_test_server(*, address, use_ssl=False, server_cls, server_ssl_cls):

    def app(environ, start_response):
        status = '200 OK'
        headers = [('Content-type', 'text/plain')]
        start_response(status, headers)
        return [b'Test message']

    # Запустить тестовый WSGI-сервер в отдельном потоке, чтобы не
    # мешать обработке событий в главном потоке
    server_class = server_ssl_cls if use_ssl else server_cls
    httpd = server_class(address, SilentWSGIRequestHandler)
    httpd.set_app(app)
    httpd.address = httpd.server_address
    server_thread = threading.Thread(
        target=lambda: httpd.serve_forever(poll_interval=0.05))
    server_thread.start()
    try:
        yield httpd
    finally:
        httpd.shutdown()
        httpd.server_close()
        server_thread.join()


if hasattr(socket, 'AF_UNIX'):

    class UnixHTTPServer(socketserver.UnixStreamServer, HTTPServer):

        def server_bind(self):
            socketserver.UnixStreamServer.server_bind(self)
            self.server_name = '127.0.0.1'
            self.server_port = 80


    class UnixWSGIServer(UnixHTTPServer, WSGIServer):

        request_timeout = 2

        def server_bind(self):
            UnixHTTPServer.server_bind(self)
            self.setup_environ()

        def get_request(self):
            request, client_addr = super().get_request()
            request.settimeout(self.request_timeout)
            # Код в stdlib ожидает, что get_request
            # вернет сокет и кортеж (host, port).
            # Однако это неверно для сокетов UNIX,
            # так как второе возвращаемое значение будет путем;
            # поэтому мы возвращаем некоторые фиктивные данные, достаточные
            # чтобы запустить тесты
            return request, ('127.0.0.1', '')


    class SilentUnixWSGIServer(UnixWSGIServer):

        def handle_error(self, request, client_address):
            pass


    class UnixSSLWSGIServer(SSLWSGIServerMixin, SilentUnixWSGIServer):
        pass


    def gen_unix_socket_path():
        with tempfile.NamedTemporaryFile() as file:
            return file.name


    @contextlib.contextmanager
    def unix_socket_path():
        path = gen_unix_socket_path()
        try:
            yield path
        finally:
            try:
                os.unlink(path)
            except OSError:
                pass


    @contextlib.contextmanager
    def run_test_unix_server(*, use_ssl=False):
        with unix_socket_path() as path:
            yield from _run_test_server(address=path, use_ssl=use_ssl,
                                        server_cls=SilentUnixWSGIServer,
                                        server_ssl_cls=UnixSSLWSGIServer)


@contextlib.contextmanager
def run_test_server(*, host='127.0.0.1', port=0, use_ssl=False):
    yield from _run_test_server(address=(host, port), use_ssl=use_ssl,
                                server_cls=SilentWSGIServer,
                                server_ssl_cls=SSLWSGIServer)


def make_test_protocol(base):
    dct = {}
    for name in dir(base):
        if name.startswith('__') and name.endswith('__'):
            # пропустить магические имена
            continue
        dct[name] = MockCallback(return_value=None)
    return type('TestProtocol', (base,) + base.__bases__, dct)()


class TestSelector(selectors.BaseSelector):

    def __init__(self):
        self.keys = {}

    def register(self, fileobj, events, data=None):
        key = selectors.SelectorKey(fileobj, 0, events, data)
        self.keys[fileobj] = key
        return key

    def unregister(self, fileobj):
        return self.keys.pop(fileobj)

    def select(self, timeout):
        return []

    def get_map(self):
        return self.keys


class TestLoop(base_events.BaseEventLoop):
    """Цикл для модульных тестов.

    Он управляет самим временем напрямую.
    Если что-то запланировано к выполнению позже, то
    на следующей итерации цикла после выполнения всех готовых обработчиков
    вызывается генератор, переданный в __init__.

    Генератор должен выглядеть так:

        def gen():
            ...
            when = yield ...
            ... = yield time_advance

    Значение, возвращаемое yield, является абсолютным временем следующего запланированного обработчика.
    Значение, передаваемое в yield, — это сдвиг времени для перемещения времени цикла вперёд.
    """

    def __init__(self, gen=None):
        super().__init__()

        if gen is None:
            def gen():
                yield
            self._check_on_close = False
        else:
            self._check_on_close = True

        self._gen = gen()
        next(self._gen)
        self._time = 0
        self._clock_resolution = 1e-9
        self._timers = []
        self._selector = TestSelector()

        self.readers = {}
        self.writers = {}
        self.reset_counters()

        self._transports = weakref.WeakValueDictionary()

    def time(self):
        return self._time

    def advance_time(self, advance):
        """Переместить тестовое время вперёд."""
        if advance:
            self._time += advance

    def close(self):
        super().close()
        if self._check_on_close:
            try:
                self._gen.send(0)
            except StopIteration:
                pass
            else:  # pragma: no cover
                raise AssertionError("Time generator is not finished")

    def _add_reader(self, fd, callback, *args):
        self.readers[fd] = events.Handle(callback, args, self)

    def _remove_reader(self, fd):
        self.remove_reader_count[fd] += 1
        if fd in self.readers:
            del self.readers[fd]
            return True
        else:
            return False

    def assert_reader(self, fd, callback, *args):
        if fd not in self.readers:
            raise AssertionError(f'fd {fd} is not registered')
        handle = self.readers[fd]
        if handle._callback != callback:
            raise AssertionError(
                f'unexpected callback: {handle._callback} != {callback}')
        if handle._args != args:
            raise AssertionError(
                f'unexpected callback args: {handle._args} != {args}')

    def assert_no_reader(self, fd):
        if fd in self.readers:
            raise AssertionError(f'fd {fd} is registered')

    def _add_writer(self, fd, callback, *args):
        self.writers[fd] = events.Handle(callback, args, self)

    def _remove_writer(self, fd):
        self.remove_writer_count[fd] += 1
        if fd in self.writers:
            del self.writers[fd]
            return True
        else:
            return False

    def assert_writer(self, fd, callback, *args):
        assert fd in self.writers, 'fd {} is not registered'.format(fd)
        handle = self.writers[fd]
        assert handle._callback == callback, '{!r} != {!r}'.format(
            handle._callback, callback)
        assert handle._args == args, '{!r} != {!r}'.format(
            handle._args, args)

    def _ensure_fd_no_transport(self, fd):
        try:
            transport = self._transports[fd]
        except KeyError:
            pass
        else:
            raise RuntimeError(
                'File descriptor {!r} is used by transport {!r}'.format(
                    fd, transport))

    def add_reader(self, fd, callback, *args):
        """Добавить колбэк чтения."""
        self._ensure_fd_no_transport(fd)
        return self._add_reader(fd, callback, *args)

    def remove_reader(self, fd):
        """Удалить колбэк чтения."""
        self._ensure_fd_no_transport(fd)
        return self._remove_reader(fd)

    def add_writer(self, fd, callback, *args):
        """Добавить колбэк записи.."""
        self._ensure_fd_no_transport(fd)
        return self._add_writer(fd, callback, *args)

    def remove_writer(self, fd):
        """Удалить колбэк записи."""
        self._ensure_fd_no_transport(fd)
        return self._remove_writer(fd)

    def reset_counters(self):
        self.remove_reader_count = collections.defaultdict(int)
        self.remove_writer_count = collections.defaultdict(int)

    def _run_once(self):
        super()._run_once()
        for when in self._timers:
            advance = self._gen.send(when)
            self.advance_time(advance)
        self._timers = []

    def call_at(self, when, callback, *args):
        self._timers.append(when)
        return super().call_at(when, callback, *args)

    def _process_events(self, event_list):
        return

    def _write_to_self(self):
        pass


def MockCallback(**kwargs):
    return mock.Mock(spec=['__call__'], **kwargs)


class MockPattern(str):
    """Строка на основе регулярного выражения с нечётким __eq__.

    Используйте этот вспомогательный класс с 'mock.assert_called_with' или в любом месте,
    где нужно сравнение строк по регулярному выражению.

    Например:
       mock_call.assert_called_with(MockPattern('spam.*ham'))
    """
    def __eq__(self, other):
        return bool(re.search(str(self), other, re.S))


def get_function_source(func):
    source = events._get_function_source(func)
    if source is None:
        raise ValueError("unable to get the source of %r" % (func,))
    return source


class TestCase(unittest.TestCase):
    @staticmethod
    def close_loop(loop):
        executor = loop._default_executor
        if executor is not None:
            executor.shutdown(wait=True)
        loop.close()

    def set_event_loop(self, loop, *, cleanup=True):
        assert loop is not None
        # убедиться, что цикл событий передаётся явно в asyncio
        events.set_event_loop(None)
        if cleanup:
            self.addCleanup(self.close_loop, loop)

    def new_test_loop(self, gen=None):
        loop = TestLoop(gen)
        self.set_event_loop(loop)
        return loop

    def unpatch_get_running_loop(self):
        events._get_running_loop = self._get_running_loop

    def setUp(self):
        self._get_running_loop = events._get_running_loop
        events._get_running_loop = lambda: None
        self._thread_cleanup = support.threading_setup()

    def tearDown(self):
        self.unpatch_get_running_loop()

        events.set_event_loop(None)

        # Обнаружить ошибку CPython #23353: убедиться, что yield/yield-from не используется
        # в блоке except генератора
        self.assertEqual(sys.exc_info(), (None, None, None))

        self.doCleanups()
        support.threading_cleanup(*self._thread_cleanup)
        support.reap_children()

    if not compat.PY34:
        # Совместимость с Python 3.3
        def subTest(self, *args, **kwargs):
            class EmptyCM:
                def __enter__(self):
                    pass
                def __exit__(self, *exc):
                    pass
            return EmptyCM()


@contextlib.contextmanager
def disable_logger():
    """Менеджер контекста для отключения регистратора asyncio.

    Например, его можно использовать для игнорирования предупреждений в режиме отладки.
    """
    old_level = logger.level
    try:
        logger.setLevel(logging.CRITICAL+1)
        yield
    finally:
        logger.setLevel(old_level)


def mock_nonblocking_socket(proto=socket.IPPROTO_TCP, type=socket.SOCK_STREAM,
                            family=socket.AF_INET):
    """Создать имитацию неблокирующего сокета."""
    sock = mock.MagicMock(socket.socket)
    sock.proto = proto
    sock.type = type
    sock.family = family
    sock.gettimeout.return_value = 0.0
    return sock


def force_legacy_ssl_support():
    return mock.patch('asyncio.sslproto._is_sslproto_available',
                      return_value=False)