nim-libp2p/tests/pubsub/utils.nim
2020-07-28 12:06:16 +09:00

47 lines
1.6 KiB
Nim

import random, options
import chronos
import ../../libp2p/standard_setup
import ../../libp2p/protocols/pubsub/gossipsub
export standard_setup
randomize()
proc generateNodes*(num: Natural, gossip: bool = false): seq[Switch] =
for i in 0..<num:
var switch = newStandardSwitch(gossip = gossip)
if gossip:
var gossipSub = GossipSub(switch.pubSub.get())
gossipSub.parameters.floodPublish = false
result.add(switch)
proc subscribeNodes*(nodes: seq[Switch]): Future[seq[Future[void]]] {.async.} =
for dialer in nodes:
for node in nodes:
if dialer.peerInfo.peerId != node.peerInfo.peerId:
await dialer.connect(node.peerInfo)
result.add(dialer.subscribePeer(node.peerInfo))
proc subscribeSparseNodes*(nodes: seq[Switch], degree: int = 2): Future[seq[Future[void]]] {.async.} =
if nodes.len < degree:
raise (ref CatchableError)(msg: "nodes count needs to be greater or equal to degree!")
for i, dialer in nodes:
if (i mod degree) != 0:
continue
for node in nodes:
if dialer.peerInfo.peerId != node.peerInfo.peerId:
await dialer.connect(node.peerInfo)
result.add(dialer.subscribePeer(node.peerInfo))
proc subscribeRandom*(nodes: seq[Switch]): Future[seq[Future[void]]] {.async.} =
for dialer in nodes:
var dialed: seq[string]
while dialed.len < nodes.len - 1:
let node = sample(nodes)
if node.peerInfo.id notin dialed:
if dialer.peerInfo.id != node.peerInfo.id:
await dialer.connect(node.peerInfo)
result.add(dialer.subscribePeer(node.peerInfo))
dialed.add(node.peerInfo.id)