import { PeerId } from "@libp2p/interface"; import { CONNECTION_LOCKED_TAG, IConnectionManager, Libp2p, Protocols } from "@waku/interfaces"; import { expect } from "chai"; import sinon from "sinon"; import { PeerManager, PeerManagerEventNames } from "./peer_manager.js"; describe("PeerManager", () => { let libp2p: Libp2p; let peerManager: PeerManager; let connectionManager: IConnectionManager; let peers: any[]; let mockConnections: any[]; const TEST_PUBSUB_TOPIC = "/test/1/waku-light-push/utf8"; const TEST_PROTOCOL = Protocols.LightPush; const clearPeerState = (): void => { (peerManager as any).lockedPeers.clear(); (peerManager as any).unlockedPeers.clear(); }; const createPeerManagerWithConfig = (numPeersToUse: number): PeerManager => { return new PeerManager({ libp2p, connectionManager: connectionManager as any, config: { numPeersToUse } }); }; const getPeersForTest = async (): Promise => { return await peerManager.getPeers({ protocol: TEST_PROTOCOL, pubsubTopic: TEST_PUBSUB_TOPIC }); }; const skipIfNoPeers = (result: PeerId[] | null): boolean => { if (!result || result.length === 0) { return true; } return false; }; beforeEach(() => { peers = [ { id: makePeerId("peer-1"), protocols: [Protocols.LightPush, Protocols.Filter, Protocols.Store] }, { id: makePeerId("peer-2"), protocols: [Protocols.LightPush, Protocols.Filter, Protocols.Store] }, { id: makePeerId("peer-3"), protocols: [Protocols.LightPush, Protocols.Filter, Protocols.Store] } ]; mockConnections = [ { remotePeer: makePeerId("peer-1"), tags: [] as string[] }, { remotePeer: makePeerId("peer-2"), tags: [] as string[] }, { remotePeer: makePeerId("peer-3"), tags: [] as string[] } ]; libp2p = mockLibp2p(mockConnections); connectionManager = { pubsubTopics: [TEST_PUBSUB_TOPIC], getConnectedPeers: async () => peers, getPeers: async () => peers, isPeerOnTopic: async (_id: PeerId, _topic: string) => true } as unknown as IConnectionManager; peerManager = new PeerManager({ libp2p, connectionManager: connectionManager as any }); clearPeerState(); (peerManager as any).isPeerAvailableForUse = () => true; }); afterEach(() => { peerManager.stop(); sinon.restore(); }); it("should initialize with default number of peers", () => { expect(peerManager["numPeersToUse"]).to.equal(2); }); it("should initialize with custom number of peers", () => { peerManager = createPeerManagerWithConfig(3); expect(peerManager["numPeersToUse"]).to.equal(3); }); it("should return available peers with correct protocol and pubsub topic", async () => { clearPeerState(); const result = await getPeersForTest(); if (skipIfNoPeers(result)) return; expect(result[0].toString()).to.equal("peer-1"); }); it("should lock peers when selected", async () => { clearPeerState(); const result = await getPeersForTest(); if (skipIfNoPeers(result)) return; expect((peerManager as any).lockedPeers.size).to.be.greaterThan(0); }); it("should unlock peer and allow reuse after renewPeer", async () => { clearPeerState(); const ids = await getPeersForTest(); if (skipIfNoPeers(ids)) return; const peerId = ids[0]; await peerManager.renewPeer(peerId, { protocol: TEST_PROTOCOL, pubsubTopic: TEST_PUBSUB_TOPIC }); expect((peerManager as any).lockedPeers.has(peerId.toString())).to.be.false; expect((peerManager as any).unlockedPeers.has(peerId.toString())).to.be .true; }); it("should not return locked peers if enough unlocked are available", async () => { clearPeerState(); const ids = await getPeersForTest(); if (skipIfNoPeers(ids)) return; (peerManager as any).lockedPeers.add(ids[0].toString()); const result = await getPeersForTest(); if (skipIfNoPeers(result)) return; expect(result).to.not.include(ids[0]); }); it("should dispatch connect and disconnect events", () => { const connectSpy = sinon.spy(); const disconnectSpy = sinon.spy(); peerManager.events.addEventListener( PeerManagerEventNames.Connect, connectSpy ); peerManager.events.addEventListener( PeerManagerEventNames.Disconnect, disconnectSpy ); peerManager["dispatchFilterPeerConnect"](peers[0].id); peerManager["dispatchFilterPeerDisconnect"](peers[0].id); expect(connectSpy.calledOnce).to.be.true; expect(disconnectSpy.calledOnce).to.be.true; }); it("should handle onConnected and onDisconnected", async () => { const peerId = peers[0].id; sinon.stub(peerManager, "isPeerOnPubsub" as any).resolves(true); await (peerManager as any).onConnected({ detail: { peerId, protocols: [Protocols.Filter] } }); await (peerManager as any).onDisconnected({ detail: peerId }); expect(true).to.be.true; }); it("should register libp2p event listeners when start is called", () => { const addEventListenerSpy = libp2p.addEventListener as sinon.SinonSpy; peerManager.start(); expect(addEventListenerSpy.calledWith("peer:identify")).to.be.true; expect(addEventListenerSpy.calledWith("peer:disconnect")).to.be.true; }); it("should unregister libp2p event listeners when stop is called", () => { const removeEventListenerSpy = libp2p.removeEventListener as sinon.SinonSpy; peerManager.stop(); expect(removeEventListenerSpy.calledWith("peer:identify")).to.be.true; expect(removeEventListenerSpy.calledWith("peer:disconnect")).to.be.true; }); it("should return only peers supporting the requested protocol and pubsub topic", async () => { peers[0].protocols = [Protocols.LightPush]; peers[1].protocols = [Protocols.Filter]; peers[2].protocols = [Protocols.Store]; (peerManager as any).isPeerAvailableForUse = () => true; const result = await getPeersForTest(); if (skipIfNoPeers(result)) return; expect(result.length).to.equal(1); expect(result[0].toString()).to.equal("peer-1"); }); it("should return exactly numPeersToUse peers when enough are available", async () => { peerManager = createPeerManagerWithConfig(2); (peerManager as any).isPeerAvailableForUse = () => true; const result = await getPeersForTest(); if (skipIfNoPeers(result)) return; expect(result.length).to.equal(2); }); it("should respect custom numPeersToUse configuration", async () => { peerManager = createPeerManagerWithConfig(1); (peerManager as any).isPeerAvailableForUse = () => true; const result = await getPeersForTest(); if (skipIfNoPeers(result)) return; expect(result.length).to.equal(1); }); it("should not return the same peer twice in consecutive getPeers calls without renew", async () => { (peerManager as any).isPeerAvailableForUse = () => true; const first = await getPeersForTest(); const second = await getPeersForTest(); expect(second.some((id: PeerId) => first.includes(id))).to.be.false; }); it("should allow a peer to be returned again after renewPeer is called", async () => { (peerManager as any).isPeerAvailableForUse = () => true; const first = await getPeersForTest(); if (skipIfNoPeers(first)) return; await peerManager.renewPeer(first[0], { protocol: TEST_PROTOCOL, pubsubTopic: TEST_PUBSUB_TOPIC }); const second = await getPeersForTest(); if (skipIfNoPeers(second)) return; expect(second).to.include(first[0]); }); it("should handle renewPeer for a non-existent or disconnected peer gracefully", async () => { const fakePeerId = { toString: () => "not-exist", equals: () => false } as any; await peerManager.renewPeer(fakePeerId, { protocol: TEST_PROTOCOL, pubsubTopic: TEST_PUBSUB_TOPIC }); expect(true).to.be.true; }); it("should add CONNECTION_LOCKED_TAG to peer connections when locking", async () => { clearPeerState(); const result = await getPeersForTest(); if (skipIfNoPeers(result)) return; const peerId = result[0]; const connection = mockConnections.find((c) => c.remotePeer.equals(peerId)); expect(connection).to.exist; expect(connection.tags).to.include(CONNECTION_LOCKED_TAG); }); it("should remove CONNECTION_LOCKED_TAG from peer connections when unlocking", async () => { clearPeerState(); const result = await getPeersForTest(); if (skipIfNoPeers(result)) return; const peerId = result[0]; await peerManager.renewPeer(peerId, { protocol: TEST_PROTOCOL, pubsubTopic: TEST_PUBSUB_TOPIC }); const connection = mockConnections.find((c) => c.remotePeer.equals(peerId)); expect(connection).to.exist; expect(connection.tags).to.not.include(CONNECTION_LOCKED_TAG); }); it("should not modify tags of connections for different peers", async () => { clearPeerState(); const result = await getPeersForTest(); if (skipIfNoPeers(result)) return; const lockedPeerId = result[0]; const otherPeerId = peers.find((p) => !p.id.equals(lockedPeerId))?.id; if (!otherPeerId) return; const otherConnection = mockConnections.find((c) => c.remotePeer.equals(otherPeerId) ); expect(otherConnection).to.exist; expect(otherConnection.tags).to.not.include(CONNECTION_LOCKED_TAG); }); }); function mockLibp2p(connections: any[]): Libp2p { return { getConnections: sinon.stub().returns(connections), getPeers: sinon .stub() .returns([ { toString: () => "peer-1" }, { toString: () => "peer-2" }, { toString: () => "peer-3" } ]), peerStore: { get: sinon.stub().callsFake((peerId: PeerId) => Promise.resolve({ id: peerId, protocols: [Protocols.LightPush, Protocols.Filter, Protocols.Store] }) ) }, dispatchEvent: sinon.spy(), addEventListener: sinon.spy(), removeEventListener: sinon.spy() } as unknown as Libp2p; } function makePeerId(id: string): PeerId { return { toString: () => id, equals: (other: any) => other && other.toString && other.toString() === id } as PeerId; }