Add WebSocketMockServer and tests for WebSocketEngine

build-logic/src/main/kotlin/Mpp.kt

@@ -141,10 +141,11 @@ fun Project.configureMpp(
  * Current Graph is something like so
  *
  * graph TB
- * commonMain --> jvmMain
- * commonMain --> appleMain
+ * commonMain --> concurrentMain
  * commonMain --> linuxMain
  * commonMain --> jsMain
+ * concurrentMain --> jvmMain
+ * concurrentMain --> appleMain
  * appleMain --> macosX64
  * appleMain --> macosArm64
  * appleMain --> iosArm64
@@ -177,12 +178,22 @@ private fun KotlinMultiplatformExtension.configureSourceSetGraph() {
   val hasAppleTarget = targets.any {
     it is KotlinNativeTarget && it.konanTarget.family in setOf(Family.IOS, Family.OSX, Family.WATCHOS, Family.TVOS)
   }
+
+  val concurrentMain = sourceSets.create("concurrentMain")
+  val concurrentTest = sourceSets.create("concurrentTest")
+
+  concurrentMain.dependsOn(sourceSets.getByName("commonMain"))
+  concurrentTest.dependsOn(sourceSets.getByName("commonTest"))
+
+  sourceSets.findByName("jvmMain")?.dependsOn(concurrentMain)
+  sourceSets.findByName("jvmTest")?.dependsOn(concurrentTest)
+
   if (hasAppleTarget) {
     val appleMain = sourceSets.create("appleMain")
     val appleTest = sourceSets.create("appleTest")
 
-    appleMain.dependsOn(sourceSets.getByName("commonMain"))
-    appleTest.dependsOn(sourceSets.getByName("commonTest"))
+    appleMain.dependsOn(concurrentMain)
+    appleTest.dependsOn(concurrentTest)
 
     allAppleTargets.forEach {
       sourceSets.findByName("${it}Main")?.dependsOn(appleMain)
@@ -240,4 +251,4 @@ fun Project.registerJavaCodegenTestTask() {
   }
 
   tasks.named("build").dependsOn(task)
-}
+}

gradle/libraries.toml

@@ -150,6 +150,9 @@ ktor-client-okhttp = { group = "io.ktor", name = "ktor-client-okhttp", version.r
 ktor-client-darwin = { group = "io.ktor", name = "ktor-client-darwin", version.ref = "ktor" }
 ktor-client-js = { group = "io.ktor", name = "ktor-client-js", version.ref = "ktor" }
 ktor-client-websockets = { group = "io.ktor", name = "ktor-client-websockets", version.ref = "ktor" }
+ktor-server-core = { group = "io.ktor", name = "ktor-server-core", version.ref = "ktor" }
+ktor-server-cio = { group = "io.ktor", name = "ktor-server-cio", version.ref = "ktor" }
+ktor-server-websockets = { group = "io.ktor", name = "ktor-server-websockets", version.ref = "ktor" }
 okhttp = { group = "com.squareup.okhttp3", name = "okhttp", version.ref = "okhttp" }
 okhttp-logging = { group = "com.squareup.okhttp3", name = "logging-interceptor", version.ref = "okhttp" }
 okhttp-mockwebserver = { group = "com.squareup.okhttp3", name = "mockwebserver", version.ref = "okhttp" }

libraries/apollo-engine-ktor/api/apollo-engine-ktor.api

@@ -1,3 +1,6 @@
 public final class com/apollographql/apollo3/network/KtorExtensionsKt {
 }
 
+public final class com/apollographql/apollo3/network/http/KtorHttpEngine_concurrentKt {
+}
+

libraries/apollo-engine-ktor/build.gradle.kts

@@ -20,20 +20,6 @@ kotlin {
       }
     }
 
-    findByName("commonTest")?.apply {
-      dependencies {
-        implementation(project(":apollo-mockserver"))
-        implementation(project(":apollo-testing-support")) {
-          because("runTest")
-          // We have a circular dependency here that creates a warning in JS
-          // w: duplicate library name: com.apollographql.apollo3:apollo-mockserver
-          // See https://youtrack.jetbrains.com/issue/KT-51110
-          // We should probably remove this circular dependency but for the time being, just use excludes
-          exclude(group =  "com.apollographql.apollo3", module = "apollo-runtime")
-        }
-      }
-    }
-
     findByName("jvmMain")?.apply {
       dependencies {
         api(libs.ktor.client.okhttp)

libraries/apollo-engine-ktor/src/commonMain/kotlin/com/apollographql/apollo3/network/http/KtorHttpEngine.kt

@@ -1,6 +1,7 @@
 package com.apollographql.apollo3.network.http
 
 import com.apollographql.apollo3.annotations.ApolloExperimental
+import com.apollographql.apollo3.annotations.ApolloInternal
 import com.apollographql.apollo3.api.http.HttpHeader
 import com.apollographql.apollo3.api.http.HttpMethod
 import com.apollographql.apollo3.api.http.HttpRequest
@@ -24,14 +25,22 @@ class KtorHttpEngine(
 
   private var disposed = false
 
+  /**
+   * @param timeoutMillis: The timeout in milliseconds used both for the connection and socket read.
+   */
   constructor(timeoutMillis: Long = 60_000) : this(timeoutMillis, timeoutMillis)
 
-  constructor(connectTimeoutMillis: Long, requestTimeoutMillis: Long) : this(
+  /**
+   * @param connectTimeoutMillis The connection timeout in milliseconds. The connection timeout is the time period in which a client should establish a connection with a server.
+   * @param readTimeoutMillis The socket read timeout in milliseconds. On JVM and Apple this maps to [HttpTimeout.HttpTimeoutCapabilityConfiguration.socketTimeoutMillis], on JS
+   * this maps to [HttpTimeout.HttpTimeoutCapabilityConfiguration.requestTimeoutMillis]
+   */
+  constructor(connectTimeoutMillis: Long, readTimeoutMillis: Long) : this(
       HttpClient {
         expectSuccess = false
         install(HttpTimeout) {
           this.connectTimeoutMillis = connectTimeoutMillis
-          this.requestTimeoutMillis = requestTimeoutMillis
+          setReadTimeout(readTimeoutMillis)
         }
       }
   )
@@ -75,3 +84,5 @@ class KtorHttpEngine(
   }
 }
 
+@ApolloInternal
+expect fun HttpTimeout.HttpTimeoutCapabilityConfiguration.setReadTimeout(readTimeoutMillis: Long)

libraries/apollo-engine-ktor/src/commonMain/kotlin/com/apollographql/apollo3/network/ws/KtorWebSocketEngine.kt

@@ -2,6 +2,7 @@ package com.apollographql.apollo3.network.ws
 
 import com.apollographql.apollo3.annotations.ApolloExperimental
 import com.apollographql.apollo3.api.http.HttpHeader
+import com.apollographql.apollo3.exception.ApolloNetworkException
 import com.apollographql.apollo3.exception.ApolloWebSocketClosedException
 import io.ktor.client.HttpClient
 import io.ktor.client.plugins.websocket.WebSockets
@@ -11,14 +12,16 @@ import io.ktor.client.request.url
 import io.ktor.http.URLBuilder
 import io.ktor.http.URLProtocol
 import io.ktor.http.Url
+import io.ktor.websocket.CloseReason
 import io.ktor.websocket.Frame
 import io.ktor.websocket.close
 import io.ktor.websocket.readText
+import kotlinx.coroutines.CancellationException
 import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.Deferred
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.SupervisorJob
 import kotlinx.coroutines.channels.Channel
-import kotlinx.coroutines.channels.ClosedReceiveChannelException
 import kotlinx.coroutines.launch
 import okio.ByteString
 
@@ -68,23 +71,24 @@ class KtorWebSocketEngine(
               val frame = sendFrameChannel.receive()
               try {
                 send(frame)
+
+                // Also close the connection if the sent frame is a close frame
+                if (frame is Frame.Close) {
+                  receiveMessageChannel.close()
+                  sendFrameChannel.close()
+                  break
+                }
               } catch (e: Exception) {
-                val closeReason = try {closeReason.await()} catch (e: Exception) {null}
-                receiveMessageChannel.close(ApolloWebSocketClosedException(code = closeReason?.code?.toInt()
-                    ?: -1, reason = closeReason?.message, cause = e))
-                sendFrameChannel.close(e)
+                handleNetworkException(e, closeReason, receiveMessageChannel, sendFrameChannel)
                 break
               }
             }
           }
           while (true) {
             when (val frame = try {
               incoming.receive()
-            } catch (e: ClosedReceiveChannelException) {
-              val closeReason = try {closeReason.await()} catch (e: Exception) {null}
-              receiveMessageChannel.close(ApolloWebSocketClosedException(code = closeReason?.code?.toInt()
-                  ?: -1, reason = closeReason?.message, cause = e))
-              sendFrameChannel.close(e)
+            } catch (e: Exception) {
+              handleNetworkException(e, closeReason, receiveMessageChannel, sendFrameChannel)
               break
             }) {
               is Frame.Text -> {
@@ -110,7 +114,7 @@ class KtorWebSocketEngine(
           }
         }
       } catch (e: Exception) {
-        receiveMessageChannel.close(e)
+        receiveMessageChannel.close(ApolloNetworkException(message = "Web socket communication error", platformCause = e))
         sendFrameChannel.close(e)
       }
     }
@@ -128,10 +132,36 @@ class KtorWebSocketEngine(
       }
 
       override fun close() {
-        sendFrameChannel.trySend(Frame.Close())
-        sendFrameChannel.close()
+        sendFrameChannel.trySend(Frame.Close(CloseReason(CLOSE_NORMAL.toShort(), "")))
       }
+    }
+  }
 
+  private suspend fun handleNetworkException(
+      e: Exception,
+      deferredCloseReason: Deferred<CloseReason?>,
+      receiveMessageChannel: Channel<String>,
+      sendFrameChannel: Channel<Frame>,
+  ) {
+    if (e is CancellationException) throw e
+    val closeReason = try {
+      deferredCloseReason.await()
+    } catch (e: Exception) {
+      null
+    }
+    val apolloException = if (closeReason != null) {
+      ApolloWebSocketClosedException(
+          code = closeReason.code.toInt(),
+          reason = closeReason.message,
+          cause = e
+      )
+    } else {
+      ApolloNetworkException(
+          message = "Web socket communication error",
+          platformCause = e
+      )
     }
+    receiveMessageChannel.close(apolloException)
+    sendFrameChannel.close(apolloException)
   }
 }

libraries/apollo-engine-ktor/src/concurrentMain/kotlin/com/apollographql/apollo3/network/http/KtorHttpEngine.concurrent.kt

@@ -0,0 +1,9 @@
+package com.apollographql.apollo3.network.http
+
+import com.apollographql.apollo3.annotations.ApolloInternal
+import io.ktor.client.plugins.HttpTimeout
+
+@ApolloInternal
+actual fun HttpTimeout.HttpTimeoutCapabilityConfiguration.setReadTimeout(readTimeoutMillis: Long) {
+  this.socketTimeoutMillis = readTimeoutMillis
+}

libraries/apollo-engine-ktor/src/jsMain/kotlin/com/apollographql/apollo3/network/http/KtorHttpEngine.js.kt

@@ -0,0 +1,10 @@
+package com.apollographql.apollo3.network.http
+
+import com.apollographql.apollo3.annotations.ApolloInternal
+import io.ktor.client.plugins.HttpTimeout
+
+@ApolloInternal
+actual fun HttpTimeout.HttpTimeoutCapabilityConfiguration.setReadTimeout(readTimeoutMillis: Long) {
+  // Cannot use socketTimeoutMillis on JS - https://youtrack.jetbrains.com/issue/KTOR-6211
+  this.requestTimeoutMillis = readTimeoutMillis
+}

libraries/apollo-mockserver/api/apollo-mockserver.api

@@ -60,3 +60,6 @@ public abstract interface class com/apollographql/apollo3/mockserver/MockServerI
 	public abstract fun url (Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
 }
 
+public final class com/apollographql/apollo3/mockserver/WebSocketMockServer_jvmKt {
+}
+

libraries/apollo-mockserver/build.gradle.kts

@@ -37,13 +37,20 @@ kotlin {
           // w: duplicate library name: com.apollographql.apollo3:apollo-mockserver
           // See https://youtrack.jetbrains.com/issue/KT-51110
           // We should probably remove this circular dependency but for the time being, just use excludes
-          exclude(group =  "com.apollographql.apollo3", module = "apollo-mockserver")
+          exclude(group = "com.apollographql.apollo3", module = "apollo-mockserver")
         }
         implementation(project(":apollo-runtime")) {
           because("We need HttpEngine for SocketTest")
         }
       }
     }
+
+    findByName("concurrentMain")?.apply {
+      dependencies {
+        implementation(libs.ktor.server.core)
+        implementation(libs.ktor.server.cio)
+        implementation(libs.ktor.server.websockets)
+      }
+    }
   }
 }
-

libraries/apollo-mockserver/src/appleMain/kotlin/com/apollographql/apollo3/mockserver/WebSocketMockServer.apple.kt

@@ -0,0 +1,7 @@
+package com.apollographql.apollo3.mockserver
+
+import com.apollographql.apollo3.annotations.ApolloInternal
+import com.apollographql.apollo3.mockserver.internal.CommonWebSocketMockServer
+
+@ApolloInternal
+actual fun WebSocketMockServer(port: Int): WebSocketMockServer = CommonWebSocketMockServer(port)

libraries/apollo-mockserver/src/commonMain/kotlin/com/apollographql/apollo3/mockserver/WebSocketMockServer.kt

@@ -0,0 +1,42 @@
+package com.apollographql.apollo3.mockserver
+
+import com.apollographql.apollo3.annotations.ApolloExperimental
+import kotlinx.coroutines.flow.Flow
+import okio.Closeable
+
+@ApolloExperimental
+interface WebSocketMockServer : Closeable {
+  @ApolloExperimental
+  sealed class WebSocketEvent {
+    @ApolloExperimental
+    class Connect(val sessionId: String, val headers: Map<String, String>) : WebSocketEvent()
+
+    @ApolloExperimental
+    class TextMessage(val sessionId: String, val text: String) : WebSocketEvent()
+
+    @ApolloExperimental
+    class BinaryMessage(val sessionId: String, val bytes: ByteArray) : WebSocketEvent()
+
+    @ApolloExperimental
+    class Close(val sessionId: String, val reasonCode: Short?, val reasonMessage: String?) : WebSocketEvent()
+
+    @ApolloExperimental
+    class Error(val sessionId: String, val cause: Throwable) : WebSocketEvent()
+  }
+
+  fun start()
+  suspend fun url(): String
+
+  val events: Flow<WebSocketEvent>
+
+  suspend fun sendText(sessionId: String, text: String)
+  suspend fun sendBinary(sessionId: String, binary: ByteArray)
+  suspend fun sendClose(sessionId: String, reasonCode: Short? = null, reasonMessage: String? = null)
+  override fun close()
+}
+
+@ApolloExperimental
+/**
+ * @param port the port to listen on. If 0, a random available port will be used.
+ */
+expect fun WebSocketMockServer(port: Int = 0): WebSocketMockServer