From 696b1518ef96c038a02edad3ca53704972e5f46e Mon Sep 17 00:00:00 2001
From: Eric Mastro <eric.mastro@gmail.com>
Date: Thu, 16 Dec 2021 16:14:49 +1100
Subject: [PATCH] replace MockConnection with TestStreamBuffer

---
 libp2p/transports/transport.nim |  8 ----
 tests/mocks/mockconnection.nim  | 72 ---------------------------------
 tests/mocks/mocktransport.nim   | 38 +++++++++++------
 3 files changed, 25 insertions(+), 93 deletions(-)
 delete mode 100644 tests/mocks/mockconnection.nim

diff --git a/libp2p/transports/transport.nim b/libp2p/transports/transport.nim
index 1806628b4..3ae1bb093 100644
--- a/libp2p/transports/transport.nim
+++ b/libp2p/transports/transport.nim
@@ -125,11 +125,3 @@ method handles*(
     .filterIt(
       it == multiCodec("p2p-circuit")
     ).len == 0
-
-method `==`*(
-    a: Transport, b: Transport): bool {.base, gcsafe.} =
-
-  return
-    a.addrs == b.addrs and
-    a.running == b.running and
-    a.listenError == b.listenError
diff --git a/tests/mocks/mockconnection.nim b/tests/mocks/mockconnection.nim
deleted file mode 100644
index bd9960944..000000000
--- a/tests/mocks/mockconnection.nim
+++ /dev/null
@@ -1,72 +0,0 @@
-## Nim-LibP2P
-## Copyright (c) 2020 Status Research & Development GmbH
-## Licensed under either of
-##  * Apache License, version 2.0, ([LICENSE-APACHE](LICENSE-APACHE))
-##  * MIT license ([LICENSE-MIT](LICENSE-MIT))
-## at your option.
-## This file may not be copied, modified, or distributed except according to
-## those terms.
-
-{.push raises: [Defect].}
-
-import std/[hashes, oids, strformat]
-import chronicles, chronos, metrics
-import ../../libp2p/[multiaddress,
-       stream/connection,
-       stream/lpstream,
-       peerinfo,
-       errors]
-import ../helpers
-
-export lpstream, peerinfo, errors
-
-logScope:
-  topics = "libp2p connection"
-
-type
-  ReadOnceMock* = proc(self: MockConnection, pbytes: pointer, nbytes: int): Future[int] {.gcsafe.}
-  WriteMock* = proc(self: MockConnection, msg: seq[byte]): Future[void] {.gcsafe.}
-
-  MockConnection* = ref object of Connection
-    readOnceMock*: ReadOnceMock
-    writeMock*: WriteMock
-
-const ReadOnceMockDefault =
-  proc(self: MockConnection, pbytes: pointer, nbytes: int): Future[int] {.async.} =
-    discard
-
-const WriteMockDefault =
-  proc(self: MockConnection, msg: seq[byte]) {.async.} =
-    discard
-
-method readOnce*(
-    self: MockConnection,
-    pbytes: pointer,
-    nbytes: int):
-    Future[int] {.async.} =
-
-  return await self.readOnceMock(self, pbytes, nbytes)
-
-method write*(self: MockConnection, msg: seq[byte]): Future[void] {.async.} =
-  await self.writeMock(self, msg)
-
-proc new*(C: typedesc[MockConnection],
-    peerId: PeerId,
-    dir: Direction = Direction.In,
-    timeout: Duration = DefaultConnectionTimeout,
-    timeoutHandler: TimeoutHandler = nil,
-    observedAddr: MultiAddress = MultiAddress(),
-    readOnceMock: ReadOnceMock = ReadOnceMockDefault,
-    writeMock: WriteMock = WriteMockDefault): MockConnection =
-
-  let conn = C(peerId: peerId,
-      dir: dir,
-      timeout: timeout,
-      timeoutHandler: timeoutHandler,
-      observedAddr: observedAddr,
-      readOnceMock: readOnceMock,
-      writeMock: writeMock)
-
-
-  conn.initStream()
-  return conn
diff --git a/tests/mocks/mocktransport.nim b/tests/mocks/mocktransport.nim
index cbb545a5a..df60855ff 100644
--- a/tests/mocks/mocktransport.nim
+++ b/tests/mocks/mocktransport.nim
@@ -1,12 +1,14 @@
 {.push raises: [Defect].}
 
-import chronos
+import chronos, chronicles
 import ../../libp2p/[multiaddress,
        stream/connection,
        transports/transport,
        upgrademngrs/upgrade]
-import ./mockconnection,
-       ../helpers
+import ../helpers
+
+logScope:
+  topics = "libp2p mocktransport"
 
 type
   StartMock* = proc(self: MockTransport, addrs: seq[MultiAddress]): Future[void] {.gcsafe.}
@@ -29,13 +31,26 @@ const StopMockDefault =
     discard
 
 const AcceptMockDefault =
-  proc(self: MockTransport): Future[Connection] =
-    let
-      fut = newFuture[Connection]("mocktransport.accept.future")
-      peerId = PeerID.init(PrivateKey.random(ECDSA, rng[]).get()).get()
-      conn = MockConnection.new(peerId)
-    fut.complete(conn)
-    return fut
+  proc(self: MockTransport): Future[Connection] {.async.} =
+    proc writeHandler(msg: seq[byte]) {.async.} =
+      return
+
+    let client = TestBufferStream.new(writeHandler)
+    proc onClose() {.async.} =
+      try:
+        let f = client.join()
+        if not f.finished: await f.cancelAndWait()
+        await allFuturesThrowing(client.close())
+
+        trace "Cleaned up client"
+
+      except CatchableError as exc:
+        let useExc {.used.} = exc
+        debug "Error cleaning up client", errMsg = exc.msg
+
+    asyncSpawn onClose()
+
+    return client
 
 proc new*(
   T: typedesc[MockTransport],
@@ -52,9 +67,6 @@ proc new*(
     acceptMock: acceptMock,
     listenError: listenError)
 
-  # if transport.listenError.isNil:
-  #   transport.listenError = ListenErrorDefault
-
   return transport
 
 method start*(self: MockTransport, addrs: seq[MultiAddress]) {.async, raises: [Defect, TransportListenError].} =