refactor userver/testsuite: use testsuite plugins

commit_hash:bfb256cc4a23796f905d3afeda5143e1b6a21df1

Author: vitja

samples/testsuite-support/tests/test_logcapture.py

@@ -1,7 +1,3 @@
-import pytest
-from pytest_userver.plugins import log_capture
-
-
 # /// [select]
 async def test_select(service_client):
     async with service_client.capture_logs(log_level='INFO') as capture:
@@ -34,33 +30,3 @@ def log_event(link, **other):
 
         call = await log_event.wait_call()
         assert call['link'] == response.headers['x-yarequestid']
-
-
-async def test_select_error(service_client):
-    async with service_client.capture_logs(log_level='INFO') as capture:
-        response = await service_client.get('/logcapture')
-        assert response.status == 200
-
-        with pytest.raises(log_capture.IncorrectUsageError):
-            capture.select(
-                text='Message to capture',
-                link=response.headers['x-yarequestid'],
-            )
-
-
-async def test_subscribe_error(service_client, mockserver):
-    async with service_client.capture_logs() as capture:
-        response = await service_client.get(
-            '/logcapture',
-            headers={'x-yatraceid': mockserver.trace_id},
-        )
-        assert response.status == 200
-
-    with pytest.raises(log_capture.IncorrectUsageError):
-
-        @capture.subscribe(
-            text='Message to capture',
-            trace_id=mockserver.trace_id,
-        )
-        def log_event(link, **other):
-            pass

testsuite/pytest_plugins/pytest_userver/client.py

@@ -20,6 +20,7 @@
 
 import aiohttp
 
+from testsuite import logcapture
 from testsuite import utils
 from testsuite.daemons import service_client
 from testsuite.utils import approx
@@ -760,7 +761,7 @@ def __init__(
         *,
         config: TestsuiteClientConfig,
         mocked_time,
-        log_capture_fixture,
+        log_capture_fixture: logcapture.CaptureServer,
         testpoint,
         testpoint_control,
         cache_invalidation_state,
@@ -905,8 +906,8 @@ async def capture_logs(
         log_level: str = 'DEBUG',
         testsuite_skip_prepare: bool = False,
     ):
-        async with self._log_capture_fixture.start_capture(
-            log_level=log_level,
+        async with self._log_capture_fixture.capture(
+            log_level=logcapture.LogLevel.from_string(log_level),
         ) as capture:
             logger.debug('Starting logcapture')
             await self._testsuite_action(
@@ -920,10 +921,9 @@ async def capture_logs(
                 await self._log_capture_fixture.wait_for_client()
                 yield capture
             finally:
-                logger.debug('Finishing logcapture')
                 await self._testsuite_action(
                     'log_capture',
-                    log_level=self._log_capture_fixture.default_log_level,
+                    log_level=self._log_capture_fixture.default_log_level.name,
                     socket_logging_duplication=False,
                     testsuite_skip_prepare=testsuite_skip_prepare,
                 )

testsuite/pytest_plugins/pytest_userver/plugins/log_capture.py

@@ -3,18 +3,12 @@
 """
 
 # pylint: disable=redefined-outer-name
-import asyncio
-import contextlib
-import enum
 import logging
-import sys
-import typing
 
 import pytest
 
-from testsuite.utils import callinfo
-from testsuite.utils import compat
-from testsuite.utils import net as net_utils
+from testsuite import logcapture
+from testsuite.logcapture import __tracebackhide__  # noqa
 
 from ..utils import tskv
 
@@ -24,165 +18,6 @@
 logger = logging.getLogger(__name__)
 
 
-class BaseError(Exception):
-    pass
-
-
-class IncorrectUsageError(BaseError):
-    pass
-
-
-class ClientConnectTimeoutError(BaseError):
-    pass
-
-
-class LogLevel(enum.Enum):
-    TRACE = 0
-    DEBUG = 1
-    INFO = 2
-    WARNING = 3
-    ERROR = 4
-    CRITICAL = 5
-    NONE = 6
-
-    @classmethod
-    def from_string(cls, level: str) -> 'LogLevel':
-        return cls[level.upper()]
-
-
-class CapturedLogs:
-    def __init__(self, *, log_level: str) -> None:
-        self._log_level = LogLevel.from_string(log_level)
-        self._logs: typing.List[tskv.TskvRow] = []
-        self._subscribers: typing.List = []
-        self._closed = False
-
-    def close(self):
-        self._closed = True
-
-    async def publish(self, row: tskv.TskvRow) -> None:
-        self._logs.append(row)
-        for query, callback in self._subscribers:
-            if _match_entry(row, query):
-                await callback(**row)
-
-    def select(self, **query) -> typing.List[tskv.TskvRow]:
-        if not self._closed:
-            raise IncorrectUsageError(
-                'select() is only supported for closed captures\nPlease move select() after context manager body',
-            )
-        level = query.get('level')
-        if level:
-            log_level = LogLevel[level]
-            if log_level.value < self._log_level.value:
-                raise IncorrectUsageError(
-                    f'Requested log level={log_level.name} is lower than service log level {self._log_level.name}',
-                )
-        result = []
-        for row in self._logs:
-            if _match_entry(row, query):
-                result.append(row)
-        return result
-
-    def subscribe(self, **query):
-        if self._closed:
-            raise IncorrectUsageError(
-                'subscribe() is not supported for closed captures\nPlease move subscribe() into context manager body',
-            )
-
-        def decorator(func):
-            decorated = callinfo.acallqueue(func)
-            self._subscribers.append((query, decorated))
-            return decorated
-
-        return decorator
-
-
-class CaptureControl:
-    def __init__(self, *, log_level: str):
-        self.default_log_level = log_level
-        self._capture: typing.Optional[CapturedLogs] = None
-        self._tasks = []
-        self._client_cond = asyncio.Condition()
-
-    async def wait_for_client(self, timeout: float = 10.0):
-        async def waiter():
-            async with self._client_cond:
-                await self._client_cond.wait_for(lambda: self._tasks)
-
-        logger.debug('Waiting for logcapture client to connect...')
-        try:
-            await asyncio.wait_for(waiter(), timeout=timeout)
-        except TimeoutError:
-            raise ClientConnectTimeoutError(
-                'Timedout while waiting for logcapture client to connect',
-            )
-
-    @compat.asynccontextmanager
-    async def start_capture(
-        self,
-        *,
-        log_level: typing.Optional[str] = None,
-        timeout: float = 10.0,
-    ):
-        if self._capture:
-            yield self._capture
-            return
-
-        if not log_level:
-            log_level = self.default_log_level
-
-        self._capture = CapturedLogs(log_level=log_level)
-        try:
-            yield self._capture
-        finally:
-            self._capture.close()
-            self._capture = None
-            if self._tasks:
-                _, pending = await asyncio.wait(self._tasks, timeout=timeout)
-                self._tasks = []
-                if pending:
-                    raise RuntimeError(
-                        'Timedout while waiting for capture task to finish',
-                    )
-
-    @compat.asynccontextmanager
-    async def start_server(self, *, sock, loop=None):
-        extra = {}
-        if sys.version_info < (3, 8):
-            if loop is None:
-                loop = asyncio.get_running_loop()
-            extra['loop'] = loop
-        server = await asyncio.start_server(
-            self._handle_client,
-            sock=sock,
-            **extra,
-        )
-        try:
-            yield server
-        finally:
-            server.close()
-            await server.wait_closed()
-
-    async def _handle_client(self, reader, writer):
-        logger.debug('logcapture client connected')
-
-        async def log_reader(capture: CapturedLogs):
-            with contextlib.closing(writer):
-                async for line in reader:
-                    row = tskv.parse_line(line.decode('utf-8'))
-                    await capture.publish(row)
-            await writer.wait_closed()
-
-        if not self._capture:
-            writer.close()
-            await writer.wait_closed()
-        else:
-            self._tasks.append(asyncio.create_task(log_reader(self._capture)))
-            async with self._client_cond:
-                self._client_cond.notify_all()
-
-
 def pytest_addoption(parser):
     group = parser.getgroup('logs-capture')
     group.addoption(
@@ -199,40 +34,27 @@ def pytest_addoption(parser):
 
 
 @pytest.fixture(scope='session')
-def userver_log_capture(_userver_capture_control, _userver_capture_server):
-    return _userver_capture_control
-
-
-@pytest.fixture(scope='session')
-def _userver_capture_control(userver_log_level):
-    return CaptureControl(log_level=userver_log_level)
-
-
-@pytest.fixture(scope='session')
-def _userver_log_capture_socket(pytestconfig):
+async def userver_log_capture(pytestconfig, userver_log_level):
     host = pytestconfig.option.logs_capture_host
     port = pytestconfig.option.logs_capture_port
     if pytestconfig.option.service_wait or pytestconfig.option.service_disable:
         port = port or DEFAULT_PORT
-    with net_utils.bind_socket(host, port) as socket:
-        yield socket
-
 
-@pytest.fixture(scope='session')
-async def _userver_capture_server(
-    _userver_capture_control: CaptureControl,
-    _userver_log_capture_socket,
-):
-    async with _userver_capture_control.start_server(
-        sock=_userver_log_capture_socket,
-    ) as server:
+    server = logcapture.CaptureServer(
+        log_level=logcapture.LogLevel.from_string(userver_log_level),
+        parse_line=_tskv_parse_line,
+    )
+    async with server.start(host=host, port=port):
         yield server
 
 
 @pytest.fixture(scope='session')
-def _userver_config_logs_capture(_userver_log_capture_socket):
+def _userver_config_logs_capture(userver_log_capture):
+    socknames = userver_log_capture.getsocknames()
+    assert socknames
+    sockname = socknames[0]
+
     def patch_config(config, _config_vars) -> None:
-        sockname = _userver_log_capture_socket.getsockname()
         logging_config = config['components_manager']['components']['logging']
         default_logger = logging_config['loggers']['default']
         # Other formats are not yet supported by log-capture.
@@ -245,12 +67,6 @@ def patch_config(config, _config_vars) -> None:
     return patch_config
 
 
-def _match_entry(row: tskv.TskvRow, query) -> bool:
-    for key, value in query.items():
-        if row.get(key) != value:
-            return False
-    return True
-
-
-def __tracebackhide__(excinfo):
-    return excinfo.errisinstance(BaseError)
+def _tskv_parse_line(rawline: bytes):
+    line = rawline.decode('utf-8')
+    return tskv.parse_line(line)