""" Test Notify and Notifee by ensuring that the proper events get called, and that the stream passed into opened_stream is correct. Note: Listen event does not get hit because MyNotifee is passed into network after network has already started listening TODO: Add tests for closed_stream, listen_close when those features are implemented in swarm """ import asyncio import enum import pytest from libp2p.network.notifee_interface import INotifee from libp2p.tools.constants import LISTEN_MADDR from libp2p.tools.factories import SwarmFactory from libp2p.tools.utils import connect_swarm class Event(enum.Enum): OpenedStream = 0 ClosedStream = 1 # Not implemented Connected = 2 Disconnected = 3 Listen = 4 ListenClose = 5 # Not implemented class MyNotifee(INotifee): def __init__(self, events): self.events = events async def opened_stream(self, network, stream): self.events.append(Event.OpenedStream) async def closed_stream(self, network, stream): # TODO: It is not implemented yet. pass async def connected(self, network, conn): self.events.append(Event.Connected) async def disconnected(self, network, conn): self.events.append(Event.Disconnected) async def listen(self, network, _multiaddr): self.events.append(Event.Listen) async def listen_close(self, network, _multiaddr): # TODO: It is not implemented yet. pass @pytest.mark.asyncio async def test_notify(is_host_secure): swarms = [SwarmFactory(is_secure=is_host_secure) for _ in range(2)] events_0_0 = [] events_1_0 = [] events_0_without_listen = [] swarms[0].register_notifee(MyNotifee(events_0_0)) swarms[1].register_notifee(MyNotifee(events_1_0)) # Listen await asyncio.gather(*[swarm.listen(LISTEN_MADDR) for swarm in swarms]) swarms[0].register_notifee(MyNotifee(events_0_without_listen)) # Connected await connect_swarm(swarms[0], swarms[1]) # OpenedStream: first await swarms[0].new_stream(swarms[1].get_peer_id()) # OpenedStream: second await swarms[0].new_stream(swarms[1].get_peer_id()) # OpenedStream: third, but different direction. await swarms[1].new_stream(swarms[0].get_peer_id()) await asyncio.sleep(0.01) # TODO: Check `ClosedStream` and `ListenClose` events after they are ready. # Disconnected await swarms[0].close_peer(swarms[1].get_peer_id()) await asyncio.sleep(0.01) # Connected again, but different direction. await connect_swarm(swarms[1], swarms[0]) await asyncio.sleep(0.01) # Disconnected again, but different direction. await swarms[1].close_peer(swarms[0].get_peer_id()) await asyncio.sleep(0.01) expected_events_without_listen = [ Event.Connected, Event.OpenedStream, Event.OpenedStream, Event.OpenedStream, Event.Disconnected, Event.Connected, Event.Disconnected, ] expected_events = [Event.Listen] + expected_events_without_listen assert events_0_0 == expected_events assert events_1_0 == expected_events assert events_0_without_listen == expected_events_without_listen # Clean up await asyncio.gather(*[swarm.close() for swarm in swarms])