From 13ca6d253dc9884cdf9bd33d190820edc31a19bc Mon Sep 17 00:00:00 2001
From: Dmitriy Ryajov <dryajov@gmail.com>
Date: Fri, 6 Sep 2019 01:13:03 -0600
Subject: [PATCH] fix: multistream tests

---
 libp2p/multistream.nim    |  4 +---
 tests/testmplex.nim       |  5 +++--
 tests/testmultistream.nim | 24 ++++++++++++++++--------
 3 files changed, 20 insertions(+), 13 deletions(-)

diff --git a/libp2p/multistream.nim b/libp2p/multistream.nim
index b14b23171..71506e84f 100644
--- a/libp2p/multistream.nim
+++ b/libp2p/multistream.nim
@@ -101,7 +101,6 @@ proc handle*(m: MultisteamSelect, conn: Connection) {.async, gcsafe.} =
   while not conn.closed:
     block main:
       var ms = cast[string](await conn.readLp())
-      echo ms
       ms.removeSuffix("\n")
       if ms.len() <= 0:
         await conn.write(m.na)
@@ -121,13 +120,12 @@ proc handle*(m: MultisteamSelect, conn: Connection) {.async, gcsafe.} =
         else:
           for h in m.handlers:
             if (not isNil(h.match) and h.match(ms)) or ms == h.proto:
-              echo h.proto
               await conn.writeLp((h.proto & "\n"))
               try:
                 await h.protocol.handler(conn, ms)
                 break main
               except Exception as exc:
-                echo exc.msg
+                echo exc.msg # TODO: Logging
           await conn.write(m.na)
 
 proc addHandler*[T: LPProtocol](m: MultisteamSelect,
diff --git a/tests/testmplex.nim b/tests/testmplex.nim
index 5be050f39..52dbd3fb2 100644
--- a/tests/testmplex.nim
+++ b/tests/testmplex.nim
@@ -3,8 +3,9 @@ import chronos, nimcrypto/utils
 import ../libp2p/connection, 
        ../libp2p/stream/lpstream, 
        ../libp2p/stream/bufferstream, 
-       ../libp2p/tcptransport, 
-       ../libp2p/transport, 
+       ../libp2p/transports/tcptransport, 
+       ../libp2p/transports/transport, 
+       ../libp2p/protocols/identify, 
        ../libp2p/multiaddress,
        ../libp2p/muxers/mplex/mplex, 
        ../libp2p/muxers/mplex/coder, 
diff --git a/tests/testmultistream.nim b/tests/testmultistream.nim
index 0cc102a7b..08d31a6ca 100644
--- a/tests/testmultistream.nim
+++ b/tests/testmultistream.nim
@@ -1,10 +1,15 @@
 import unittest, strutils, sequtils, sugar, strformat
 import chronos
-import ../libp2p/connection, ../libp2p/multistream,
-       ../libp2p/stream/lpstream, ../libp2p/connection, 
-       ../libp2p/multiaddress, ../libp2p/transport, 
-       ../libp2p/tcptransport, ../libp2p/protocol,
-       ../libp2p/crypto/crypto, ../libp2p/peerinfo,
+import ../libp2p/connection, 
+       ../libp2p/multistream,
+       ../libp2p/stream/lpstream, 
+       ../libp2p/connection, 
+       ../libp2p/multiaddress, 
+       ../libp2p/transports/transport, 
+       ../libp2p/transports/tcptransport, 
+       ../libp2p/protocols/protocol,
+       ../libp2p/crypto/crypto, 
+       ../libp2p/peerinfo,
        ../libp2p/peer
 
 ## Mock stream for select test
@@ -41,6 +46,8 @@ method readExactly*(s: TestSelectStream,
 method write*(s: TestSelectStream, msg: string, msglen = -1) 
   {.async, gcsafe.} = discard
 
+method close(s: TestSelectStream) {.async, gcsafe.} = s.closed = true
+
 proc newTestSelectStream(): TestSelectStream =
   new result
   result.step = 1
@@ -85,7 +92,7 @@ method write*(s: TestLsStream, msg: seq[byte], msglen = -1) {.async, gcsafe.} =
 method write*(s: TestLsStream, msg: string, msglen = -1)
   {.async, gcsafe.} = discard
 
-method close*(s: TestLsStream) {.async, gcsafe.} = discard
+method close(s: TestLsStream) {.async, gcsafe.} = s.closed = true
 
 proc newTestLsStream(ls: LsHandler): TestLsStream =
   new result
@@ -131,7 +138,7 @@ method write*(s: TestNaStream, msg: string, msglen = -1) {.async, gcsafe.} =
   if s.step == 4:
     await s.na(msg)
 
-method close*(s: TestNaStream) {.async, gcsafe.} = discard
+method close(s: TestNaStream) {.async, gcsafe.} = s.closed = true
 
 proc newTestNaStream(na: NaHandler): TestNaStream =
   new result
@@ -161,6 +168,7 @@ suite "Multistream select":
                        proto: string): 
                        Future[void] {.async, gcsafe.} =
         check proto == "/test/proto/1.0.0"
+        await conn.close()
 
       protocol.handler = testHandler
       ms.addHandler("/test/proto/1.0.0", protocol)
@@ -254,7 +262,7 @@ suite "Multistream select":
       let transport2: TcpTransport = newTransport(TcpTransport)
       let conn = await transport2.dial(ma)
 
-      check (await msDial.select(conn, "/test/proto/1.0.0")) == "/test/proto/1.0.0"
+      check (await msDial.select(conn, "/test/proto/1.0.0")) == true
 
       let hello = cast[string](await conn.readLp())
       result = hello == "Hello!"