import std/sequtils import std/tables import pkg/asynctest import pkg/chronos import pkg/libp2p import pkg/libp2p/errors import pkg/codex/rng import pkg/codex/chunker import pkg/codex/blocktype as bt import pkg/codex/blockexchange import ../helpers import ../examples suite "Network - Handlers": let rng = Rng.instance() seckey = PrivateKey.random(rng[]).tryGet() peerId = PeerID.init(seckey.getPublicKey().tryGet()).tryGet() chunker = RandomChunker.new(Rng.instance(), size = 1024, chunkSize = 256) var network: BlockExcNetwork networkPeer: NetworkPeer buffer: BufferStream blocks: seq[bt.Block] done: Future[void] proc getConn(): Future[Connection] {.async.} = return Connection(buffer) setup: while true: let chunk = await chunker.getBytes() if chunk.len <= 0: break blocks.add(bt.Block.new(chunk).tryGet()) done = newFuture[void]() buffer = BufferStream.new() network = BlockExcNetwork.new( switch = newStandardSwitch(), connProvider = getConn) network.setupPeer(peerId) networkPeer = network.peers[peerId] discard await networkPeer.connect() test "Want List handler": proc wantListHandler(peer: PeerID, wantList: WantList) {.gcsafe, async.} = # check that we got the correct amount of entries check wantList.entries.len == 4 for b in blocks: check b.cid in wantList.entries let entry = wantList.entries[wantList.entries.find(b.cid)] check entry.wantType == WantType.wantHave check entry.priority == 1 check entry.cancel == true check entry.sendDontHave == true done.complete() network.handlers.onWantList = wantListHandler let wantList = makeWantList( blocks.mapIt( it.cid ), 1, true, WantType.wantHave, true, true) let msg = Message(wantlist: wantList) await buffer.pushData(lenPrefix(ProtobufEncode(msg))) await done.wait(500.millis) test "Blocks Handler": proc blocksHandler(peer: PeerID, blks: seq[bt.Block]) {.gcsafe, async.} = check blks == blocks done.complete() network.handlers.onBlocks = blocksHandler let msg = Message(payload: makeBlocks(blocks)) await buffer.pushData(lenPrefix(ProtobufEncode(msg))) await done.wait(500.millis) test "Presence Handler": proc presenceHandler( peer: PeerID, precense: seq[BlockPresence]) {.gcsafe, async.} = for b in blocks: check: b.cid in precense done.complete() network.handlers.onPresence = presenceHandler let msg = Message( blockPresences: blocks.mapIt( BlockPresence( cid: it.cid.data.buffer, type: BlockPresenceType.presenceHave ))) await buffer.pushData(lenPrefix(ProtobufEncode(msg))) await done.wait(500.millis) test "Handles account messages": let account = Account(address: EthAddress.example) proc handleAccount(peer: PeerID, received: Account) {.gcsafe, async.} = check received == account done.complete() network.handlers.onAccount = handleAccount let message = Message(account: AccountMessage.init(account)) await buffer.pushData(lenPrefix(ProtobufEncode(message))) await done.wait(100.millis) test "Handles payment messages": let payment = SignedState.example proc handlePayment(peer: PeerID, received: SignedState) {.gcsafe, async.} = check received == payment done.complete() network.handlers.onPayment = handlePayment let message = Message(payment: StateChannelUpdate.init(payment)) await buffer.pushData(lenPrefix(ProtobufEncode(message))) await done.wait(100.millis) suite "Network - Senders": let chunker = RandomChunker.new(Rng.instance(), size = 1024, chunkSize = 256) var switch1, switch2: Switch network1, network2: BlockExcNetwork blocks: seq[bt.Block] done: Future[void] setup: while true: let chunk = await chunker.getBytes() if chunk.len <= 0: break blocks.add(bt.Block.new(chunk).tryGet()) done = newFuture[void]() switch1 = newStandardSwitch() switch2 = newStandardSwitch() network1 = BlockExcNetwork.new( switch = switch1) switch1.mount(network1) network2 = BlockExcNetwork.new( switch = switch2) switch2.mount(network2) await switch1.start() await switch2.start() await switch1.connect( switch2.peerInfo.peerId, switch2.peerInfo.addrs) teardown: await allFuturesThrowing( switch1.stop(), switch2.stop()) test "Send want list": proc wantListHandler(peer: PeerID, wantList: WantList) {.gcsafe, async.} = # check that we got the correct amount of entries check wantList.entries.len == 4 for b in blocks: check b.cid in wantList.entries let entry = wantList.entries[wantList.entries.find(b.cid)] check entry.wantType == WantType.wantHave check entry.priority == 1 check entry.cancel == true check entry.sendDontHave == true done.complete() network2.handlers.onWantList = wantListHandler await network1.sendWantList( switch2.peerInfo.peerId, blocks.mapIt( it.cid ), 1, true, WantType.wantHave, true, true) await done.wait(500.millis) test "send blocks": proc blocksHandler(peer: PeerID, blks: seq[bt.Block]) {.gcsafe, async.} = check blks == blocks done.complete() network2.handlers.onBlocks = blocksHandler await network1.sendBlocks( switch2.peerInfo.peerId, blocks) await done.wait(500.millis) test "send presence": proc presenceHandler( peer: PeerID, precense: seq[BlockPresence]) {.gcsafe, async.} = for b in blocks: check: b.cid in precense done.complete() network2.handlers.onPresence = presenceHandler await network1.sendBlockPresence( switch2.peerInfo.peerId, blocks.mapIt( BlockPresence( cid: it.cid.data.buffer, type: BlockPresenceType.presenceHave ))) await done.wait(500.millis) test "send account": let account = Account(address: EthAddress.example) proc handleAccount(peer: PeerID, received: Account) {.gcsafe, async.} = check received == account done.complete() network2.handlers.onAccount = handleAccount await network1.sendAccount(switch2.peerInfo.peerId, account) await done.wait(500.millis) test "send payment": let payment = SignedState.example proc handlePayment(peer: PeerID, received: SignedState) {.gcsafe, async.} = check received == payment done.complete() network2.handlers.onPayment = handlePayment await network1.sendPayment(switch2.peerInfo.peerId, payment) await done.wait(500.millis) suite "Network - Test Limits": var switch1, switch2: Switch network1, network2: BlockExcNetwork blocks: seq[bt.Block] done: Future[void] setup: done = newFuture[void]() switch1 = newStandardSwitch() switch2 = newStandardSwitch() network1 = BlockExcNetwork.new( switch = switch1, maxInflight = 0) switch1.mount(network1) network2 = BlockExcNetwork.new( switch = switch2) switch2.mount(network2) await switch1.start() await switch2.start() await switch1.connect( switch2.peerInfo.peerId, switch2.peerInfo.addrs) teardown: await allFuturesThrowing( switch1.stop(), switch2.stop()) test "Concurrent Sends": let account = Account(address: EthAddress.example) network2.handlers.onAccount = proc(peer: PeerID, received: Account) {.gcsafe, async.} = check false let fut = network1.send( switch2.peerInfo.peerId, Message(account: AccountMessage.init(account))) await sleepAsync(100.millis) check not fut.finished