py-libp2p/tests_interop/conftest.py

import asyncio
import sys
from typing import Union

from p2pclient.datastructures import StreamInfo
import pexpect
import pytest

from libp2p.io.abc import ReadWriteCloser
from tests.configs import LISTEN_MADDR
from tests.factories import (
    FloodsubFactory,
    GossipsubFactory,
    HostFactory,
    PubsubFactory,
)
from tests.pubsub.configs import GOSSIPSUB_PARAMS

from .daemon import Daemon, make_p2pd
from .utils import connect


@pytest.fixture
def is_host_secure():
    return False


@pytest.fixture
def num_hosts():
    return 3


@pytest.fixture
async def hosts(num_hosts, is_host_secure):
    _hosts = HostFactory.create_batch(num_hosts, is_secure=is_host_secure)
    await asyncio.gather(
        *[_host.get_network().listen(LISTEN_MADDR) for _host in _hosts]
    )
    try:
        yield _hosts
    finally:
        # TODO: It's possible that `close` raises exceptions currently,
        #   due to the connection reset things. Though we don't care much about that when
        #   cleaning up the tasks, it is probably better to handle the exceptions properly.
        await asyncio.gather(
            *[_host.close() for _host in _hosts], return_exceptions=True
        )


@pytest.fixture
def proc_factory():
    procs = []

    def call_proc(cmd, args, logfile=None, encoding=None):
        if logfile is None:
            logfile = sys.stdout
        if encoding is None:
            encoding = "utf-8"
        proc = pexpect.spawn(cmd, args, logfile=logfile, encoding=encoding)
        procs.append(proc)
        return proc

    try:
        yield call_proc
    finally:
        for proc in procs:
            proc.close()


@pytest.fixture
def num_p2pds():
    return 1


@pytest.fixture
def is_gossipsub():
    return True


@pytest.fixture
async def p2pds(num_p2pds, is_host_secure, is_gossipsub, unused_tcp_port_factory):
    p2pds: Union[Daemon, Exception] = await asyncio.gather(
        *[
            make_p2pd(
                unused_tcp_port_factory(),
                unused_tcp_port_factory(),
                is_host_secure,
                is_gossipsub=is_gossipsub,
            )
            for _ in range(num_p2pds)
        ],
        return_exceptions=True,
    )
    p2pds_succeeded = tuple(p2pd for p2pd in p2pds if isinstance(p2pd, Daemon))
    if len(p2pds_succeeded) != len(p2pds):
        # Not all succeeded. Close the succeeded ones and print the failed ones(exceptions).
        await asyncio.gather(*[p2pd.close() for p2pd in p2pds_succeeded])
        exceptions = tuple(p2pd for p2pd in p2pds if isinstance(p2pd, Exception))
        raise Exception(f"not all p2pds succeed: first exception={exceptions[0]}")
    try:
        yield p2pds
    finally:
        await asyncio.gather(*[p2pd.close() for p2pd in p2pds])


@pytest.fixture
def pubsubs(num_hosts, hosts, is_gossipsub):
    if is_gossipsub:
        routers = GossipsubFactory.create_batch(num_hosts, **GOSSIPSUB_PARAMS._asdict())
    else:
        routers = FloodsubFactory.create_batch(num_hosts)
    _pubsubs = tuple(
        PubsubFactory(host=host, router=router) for host, router in zip(hosts, routers)
    )
    yield _pubsubs
    # TODO: Clean up


class DaemonStream(ReadWriteCloser):
    stream_info: StreamInfo
    reader: asyncio.StreamReader
    writer: asyncio.StreamWriter

    def __init__(
        self,
        stream_info: StreamInfo,
        reader: asyncio.StreamReader,
        writer: asyncio.StreamWriter,
    ) -> None:
        self.stream_info = stream_info
        self.reader = reader
        self.writer = writer

    async def close(self) -> None:
        self.writer.close()
        await self.writer.wait_closed()

    async def read(self, n: int = -1) -> bytes:
        return await self.reader.read(n)

    async def write(self, data: bytes) -> int:
        return self.writer.write(data)


@pytest.fixture
async def is_to_fail_daemon_stream():
    return False


@pytest.fixture
async def py_to_daemon_stream_pair(hosts, p2pds, is_to_fail_daemon_stream):
    assert len(hosts) >= 1
    assert len(p2pds) >= 1
    host = hosts[0]
    p2pd = p2pds[0]
    protocol_id = "/protocol/id/123"
    stream_py = None
    stream_daemon = None
    event_stream_handled = asyncio.Event()
    await connect(host, p2pd)

    async def daemon_stream_handler(stream_info, reader, writer):
        nonlocal stream_daemon
        stream_daemon = DaemonStream(stream_info, reader, writer)
        event_stream_handled.set()

    await p2pd.control.stream_handler(protocol_id, daemon_stream_handler)
    # Sleep for a while to wait for the handler being registered.
    await asyncio.sleep(0.01)

    if is_to_fail_daemon_stream:
        # FIXME: This is a workaround to make daemon reset the stream.
        #   We intentionally close the listener on the python side, it makes the connection from
        #   daemon to us fail, and therefore the daemon resets the opened stream on their side.
        #   Reference: https://github.com/libp2p/go-libp2p-daemon/blob/b95e77dbfcd186ccf817f51e95f73f9fd5982600/stream.go#L47-L50  # noqa: E501
        #   We need it because we want to test against `stream_py` after the remote side(daemon)
        #   is reset. This should be removed after the API `stream.reset` is exposed in daemon
        #   some day.
        listener = p2pds[0].control.control.listener
        listener.close()
        await listener.wait_closed()
    stream_py = await host.new_stream(p2pd.peer_id, [protocol_id])
    if not is_to_fail_daemon_stream:
        await event_stream_handled.wait()
    # NOTE: If `is_to_fail_daemon_stream == True`, then `stream_daemon == None`.
    yield stream_py, stream_daemon