image-match/python/py/Lib/site-packages/jupyter_lsp/tests/test_stdio.py

import asyncio
import subprocess
import sys

import pytest
from tornado.queues import Queue

from jupyter_lsp.stdio import LspStdIoReader

WRITER_TEMPLATE = """
from time import sleep

print('Content-Length: {length}')
print()

for repeat in range({repeats}):
    sleep({interval})
    print('{message}', end='')

if {add_excess}:
    print("extra", end='')

print()
"""


class CommunicatorSpawner:
    def __init__(self, tmp_path):
        self.tmp_path = tmp_path

    def spawn_writer(
        self, message: str, repeats: int = 1, interval=None, add_excess=False
    ):
        length = len(message) * repeats
        commands_file = self.tmp_path / "writer.py"
        commands_file.write_text(
            WRITER_TEMPLATE.format(
                length=length,
                repeats=repeats,
                interval=interval or 0,
                message=message,
                add_excess=add_excess,
            )
        )
        return subprocess.Popen(
            [sys.executable, "-u", str(commands_file)],
            stdout=subprocess.PIPE,
            bufsize=0,
        )


@pytest.fixture
def communicator_spawner(tmp_path):
    return CommunicatorSpawner(tmp_path)


async def join_process(process: subprocess.Popen, headstart=1, timeout=1):
    await asyncio.sleep(headstart)
    result = process.wait(timeout=timeout)
    if process.stdout:
        process.stdout.close()
    return result


@pytest.mark.parametrize(
    "message,repeats,interval,add_excess",
    [
        ["short", 1, None, False],
        ["ab" * 10_0000, 1, None, False],
        ["ab", 2, 0.01, False],
        ["ab", 45, 0.01, False],
        ["message", 2, 0.01, True],
    ],
    ids=["short", "long", "intermittent", "intensive-intermittent", "with-excess"],
)
@pytest.mark.asyncio
async def test_reader(message, repeats, interval, add_excess, communicator_spawner):
    queue = Queue()

    process = communicator_spawner.spawn_writer(
        message=message, repeats=repeats, interval=interval, add_excess=add_excess
    )
    reader = LspStdIoReader(stream=process.stdout, queue=queue)

    await asyncio.gather(join_process(process, headstart=3, timeout=1), reader.read())

    result = queue.get_nowait()
    assert result == message * repeats