test(transports): fix muxer handling (#1884)

Author: vladopajic

tests/transports/stream_tests.nim

@@ -116,7 +116,6 @@ template streamTransportTest*(
       let client = transportProvider()
       let conn = await client.dial("", server.addrs[0])
       let muxer = streamProvider(conn)
-      discard muxer.handle()
 
       let stream = await muxer.newStream()
       await stream.write(serverMessage)
@@ -263,7 +262,6 @@ template streamTransportTest*(
       let client = transportProvider()
       let conn = await client.dial(server.addrs[0])
       let muxer = streamProvider(conn)
-      discard muxer.handle()
 
       # Send incomplete messages (will block)
       const incompleteClientMessage = clientMessage[0 ..< 10]
@@ -394,7 +392,6 @@ template streamTransportTest*(
       let client = transportProvider()
       let conn = await client.dial(server.addrs[0])
       let muxer = streamProvider(conn)
-      discard muxer.handle()
 
       var futs: seq[Future[void]]
       for i in 0 ..< numStreams:
@@ -463,11 +460,9 @@ template streamTransportTest*(
       # Accept multiple connections and handle them
       var futs: seq[Future[void]]
       for i in 0 ..< numConnections:
-        let conn = await server.accept()
-        let muxer = streamProvider(conn)
-
         # Use a proc to properly capture loop index
-        proc setupConnection(conn: Connection, muxer: Muxer, handlerIndex: int) =
+        proc setupConnection(conn: Connection, handlerIndex: int) =
+          let muxer = streamProvider(conn, false)
           muxer.streamHandler = proc(stream: Connection) {.async: (raises: []).} =
             noExceptionWithStreamClose(stream):
               # Read data in chunks with random delay
@@ -503,14 +498,14 @@ template streamTransportTest*(
 
           futs.add(startStreamHandlerAndCleanup())
 
-        setupConnection(conn, muxer, i)
+        let conn = await server.accept()
+        setupConnection(conn, i)
       await allFutures(futs)
 
     proc runClient(server: Transport, connectionId: int) {.async.} =
       let client = transportProvider()
       let conn = await client.dial(server.addrs[0])
       let muxer = streamProvider(conn)
-      discard muxer.handle()
 
       let stream = await muxer.newStream()
 

tests/transports/testquic.nim

@@ -24,7 +24,7 @@ proc quicTransProvider(): Transport {.gcsafe, raises: [].} =
   except ResultError[crypto.CryptoError]:
     raiseAssert "should not happen"
 
-proc streamProvider(conn: Connection): Muxer {.raises: [].} =
+proc streamProvider(conn: Connection, handle: bool = true): Muxer {.raises: [].} =
   try:
     return QuicMuxer.new(conn)
   except CatchableError:

tests/transports/testtcp.nim

@@ -27,8 +27,11 @@ import ./tcp_tests
 proc tcpTransProvider(): Transport =
   TcpTransport.new(upgrade = Upgrade())
 
-proc streamProvider(conn: Connection): Muxer =
-  Mplex.new(conn)
+proc streamProvider(conn: Connection, handle: bool = true): Muxer =
+  let muxer = Mplex.new(conn)
+  if handle:
+    asyncSpawn muxer.handle()
+  muxer
 
 const
   addressIP4 = "/ip4/127.0.0.1/tcp/0"

tests/transports/testtor.nim

@@ -37,8 +37,11 @@ suite "Tor transport":
   proc torTransProvider(): Transport =
     TorTransport.new(torServer, {ReuseAddr}, Upgrade())
 
-  proc streamProvider(conn: Connection): Muxer =
-    Mplex.new(conn)
+  proc streamProvider(conn: Connection, handle: bool = true): Muxer =
+    let muxer = Mplex.new(conn)
+    if handle:
+      asyncSpawn muxer.handle()
+    muxer
 
   const
     address =

tests/transports/testws.nim

@@ -83,8 +83,11 @@ proc wsSecureTransProvider(): Transport {.gcsafe, raises: [].} =
   except TLSStreamProtocolError:
     raiseAssert "should not happen"
 
-proc streamProvider(conn: Connection): Muxer =
-  Mplex.new(conn)
+proc streamProvider(conn: Connection, handle: bool = true): Muxer =
+  let muxer = Mplex.new(conn)
+  if handle:
+    asyncSpawn muxer.handle()
+  muxer
 
 const
   wsAddress = "/ip4/127.0.0.1/tcp/0/ws"

tests/transports/utils.nim

@@ -106,7 +106,8 @@ proc createTransport*(
 
 type TransportProvider* = proc(): Transport {.gcsafe, raises: [].}
 
-type StreamProvider* = proc(conn: Connection): Muxer {.gcsafe, raises: [].}
+type StreamProvider* =
+  proc(conn: Connection, handle: bool = true): Muxer {.gcsafe, raises: [].}
 
 type StreamHandler* = proc(stream: Connection) {.async: (raises: []).}
 
@@ -137,12 +138,11 @@ proc serverHandlerSingleStream*(
 ) {.async: (raises: []).} =
   try:
     let conn = await server.accept()
-    let muxer = streamProvider(conn)
+    let muxer = streamProvider(conn, false)
     muxer.streamHandler = handler
 
-    let muxerTask = muxer.handle()
+    await muxer.handle()
 
-    await muxerTask
     await muxer.close()
     await conn.close()
   except CatchableError as exc:
@@ -158,7 +158,6 @@ proc clientRunSingleStream*(
     let client = transportProvider()
     let conn = await client.dial("", server.addrs[0])
     let muxer = streamProvider(conn)
-    discard muxer.handle()
 
     let stream = await muxer.newStream()
     await handler(stream)
@@ -180,10 +179,10 @@ proc runSingleStreamScenario*(
   let serverTask =
     serverHandlerSingleStream(server, streamProvider, serverStreamHandler)
 
-  await clientRunSingleStream(
+  let clientTask = clientRunSingleStream(
     server, transportProvider, streamProvider, clientStreamHandler
   )
-  await serverTask
+  await allFutures(clientTask, serverTask)
   await server.stop()
 
 proc countTransitions*(readOrder: seq[byte]): int =