From 9418e2e52ba481a843ea144baabc4ef1768a3189 Mon Sep 17 00:00:00 2001
From: Sergey Mashkov <sergey.mashkov@jetbrains.com>
Date: Tue, 11 Dec 2018 21:36:08 +0300
Subject: [PATCH] Fix testUpgrade to not wait for a socket to be closed by
 server

---
 .../io/ktor/server/testing/EngineTestSuite.kt | 25 ++++++++++++++-----
 1 file changed, 19 insertions(+), 6 deletions(-)

diff --git a/ktor-server/ktor-server-test-host/src/io/ktor/server/testing/EngineTestSuite.kt b/ktor-server/ktor-server-test-host/src/io/ktor/server/testing/EngineTestSuite.kt
index fb676e73d2d..2c8fb5c0fe4 100644
--- a/ktor-server/ktor-server-test-host/src/io/ktor/server/testing/EngineTestSuite.kt
+++ b/ktor-server/ktor-server-test-host/src/io/ktor/server/testing/EngineTestSuite.kt
@@ -1361,6 +1361,8 @@ abstract class EngineTestSuite<TEngine : ApplicationEngine, TConfiguration : App
 
     @Test
     open fun testUpgrade() {
+        val completed = CompletableDeferred<Unit>()
+
         createAndStartServer {
             get("/up") {
                 call.respond(object : OutgoingContent.ProtocolUpgrade() {
@@ -1372,11 +1374,20 @@ abstract class EngineTestSuite<TEngine : ApplicationEngine, TConfiguration : App
 
                     override suspend fun upgrade(input: ByteReadChannel, output: ByteWriteChannel, engineContext: CoroutineContext, userContext: CoroutineContext): Job {
                         return launch(engineContext) {
-                            val bb = ByteBuffer.allocate(8)
-                            input.readFully(bb)
-                            bb.flip()
-                            output.writeFully(bb)
-                            output.close()
+                            try {
+                                val bb = ByteBuffer.allocate(8)
+                                input.readFully(bb)
+                                bb.flip()
+                                output.writeFully(bb)
+                                output.close()
+                                input.readRemaining().use {
+                                    assertEquals(0, it.remaining)
+                                }
+                                completed.complete(Unit)
+                            } catch (t: Throwable) {
+                                completed.completeExceptionally(t)
+                                throw t
+                            }
                         }
                     }
                 })
@@ -1444,7 +1455,9 @@ abstract class EngineTestSuite<TEngine : ApplicationEngine, TConfiguration : App
 
                 assertEquals(0x1122334455667788L, ch.readLong())
 
-                assertEquals(-1, ch.readAvailable(ByteArray(1)))
+                close()
+
+                completed.await()
             }
         }
     }