chore: mock websocket session [WPB-9999]

logic/src/commonMain/kotlin/com/wire/kalium/logic/feature/UserSessionScope.kt

@@ -446,6 +446,7 @@ import com.wire.kalium.network.NetworkStateObserver
 import com.wire.kalium.network.networkContainer.AuthenticatedNetworkContainer
 import com.wire.kalium.network.session.SessionManager
 import com.wire.kalium.network.utils.MockUnboundNetworkClient
+import com.wire.kalium.network.utils.MockWebSocketSession
 import com.wire.kalium.persistence.client.ClientRegistrationStorage
 import com.wire.kalium.persistence.client.ClientRegistrationStorageImpl
 import com.wire.kalium.persistence.db.GlobalDatabaseBuilder
@@ -586,6 +587,7 @@ class UserSessionScope internal constructor(
         userAgent = userAgent,
         certificatePinning = kaliumConfigs.certPinningConfig,
         mockEngine = kaliumConfigs.mockedRequests?.let { MockUnboundNetworkClient.createMockEngine(it) },
+        mockWebSocketSession = if (kaliumConfigs.mockedWebSocket) MockWebSocketSession() else null,
         kaliumLogger = userScopedLogger
     )
     private val featureSupport: FeatureSupport = FeatureSupportImpl(
@@ -916,11 +918,12 @@ class UserSessionScope internal constructor(
             kaliumFileSystem = kaliumFileSystem
         )
 
-    private val eventGatherer: EventGatherer get() = EventGathererImpl(
-        eventRepository,
-        incrementalSyncRepository,
-        userScopedLogger,
-    )
+    private val eventGatherer: EventGatherer
+        get() = EventGathererImpl(
+            eventRepository,
+            incrementalSyncRepository,
+            userScopedLogger,
+        )
 
     private val eventProcessor: EventProcessor by lazy {
         EventProcessorImpl(
@@ -1994,12 +1997,13 @@ class UserSessionScope internal constructor(
             appLockConfigHandler
         )
 
-    val team: TeamScope get() = TeamScope(
-        teamRepository = teamRepository,
-        conversationRepository = conversationRepository,
-        slowSyncRepository = slowSyncRepository,
-        selfTeamIdProvider = selfTeamId
-    )
+    val team: TeamScope
+        get() = TeamScope(
+            teamRepository = teamRepository,
+            conversationRepository = conversationRepository,
+            slowSyncRepository = slowSyncRepository,
+            selfTeamIdProvider = selfTeamId
+        )
 
     val service: ServiceScope
         get() = ServiceScope(

logic/src/commonMain/kotlin/com/wire/kalium/logic/featureFlags/KaliumConfigs.kt

@@ -43,6 +43,7 @@ data class KaliumConfigs(
     val certPinningConfig: Map<String, List<String>> = emptyMap(),
     val mockedRequests: List<TestRequestHandler>? = null,
     val mockNetworkStateObserver: NetworkStateObserver? = null,
+    val mockedWebSocket: Boolean = false,
     // Interval between attempts to advance the proteus to MLS migration
     val mlsMigrationInterval: Duration = 24.hours,
     // limit for the number of team members to fetch during slow sync

network/src/commonMain/kotlin/com/wire/kalium/network/NetworkClient.kt

@@ -36,7 +36,10 @@ import io.ktor.client.plugins.compression.ContentEncoding
 import io.ktor.client.plugins.contentnegotiation.ContentNegotiation
 import io.ktor.client.plugins.logging.LogLevel
 import io.ktor.client.plugins.websocket.WebSockets
+import io.ktor.client.plugins.websocket.webSocketSession
+import io.ktor.client.request.HttpRequestBuilder
 import io.ktor.serialization.kotlinx.json.json
+import io.ktor.websocket.WebSocketSession
 
 /**
  * Provides a [HttpClient] that has all the
@@ -102,6 +105,7 @@ internal class AuthenticatedWebSocketClient(
     private val bearerAuthProvider: BearerAuthProvider,
     private val serverConfigDTO: ServerConfigDTO,
     private val kaliumLogger: KaliumLogger,
+    private val webSocketSessionProvider: ((HttpClient, String) -> WebSocketSession)? = null
 ) {
     /**
      * Creates a disposable [HttpClient] for a single use.
@@ -123,6 +127,13 @@ internal class AuthenticatedWebSocketClient(
                 pingInterval = WEBSOCKET_PING_INTERVAL_MILLIS
             }
         }
+
+    suspend fun createWebSocketSession(clientId: String, block: HttpRequestBuilder.() -> Unit): WebSocketSession {
+        val client = createDisposableHttpClient()
+        return webSocketSessionProvider?.let {
+            return it(client, clientId)
+        } ?: client.webSocketSession(block)
+    }
 }
 
 internal fun provideBaseHttpClient(

network/src/commonMain/kotlin/com/wire/kalium/network/api/v0/authenticated/NotificationApiV0.kt

@@ -43,13 +43,12 @@ import com.wire.kalium.network.utils.deleteSensitiveItemsFromJson
 import com.wire.kalium.network.utils.mapSuccess
 import com.wire.kalium.network.utils.setWSSUrl
 import com.wire.kalium.network.utils.wrapKaliumResponse
-import io.ktor.client.plugins.websocket.DefaultClientWebSocketSession
-import io.ktor.client.plugins.websocket.webSocket
 import io.ktor.client.request.get
 import io.ktor.client.request.parameter
 import io.ktor.http.HttpStatusCode
 import io.ktor.http.Url
 import io.ktor.websocket.Frame
+import io.ktor.websocket.WebSocketSession
 import io.ktor.websocket.close
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.FlowCollector
@@ -118,19 +117,16 @@ internal open class NotificationApiV0 internal constructor(
                 //       exceptions when the backend returns 401 instead of triggering a token refresh.
                 //       This call to lastNotification will make sure that if the token is expired, it will be refreshed
                 //       before attempting to open the websocket
-                authenticatedWebSocketClient
-                    .createDisposableHttpClient()
-                    .webSocket({
-                        setWSSUrl(Url(serverLinks.webSocket), PATH_AWAIT)
-                        parameter(CLIENT_QUERY_KEY, clientId)
-                    }) {
-                        emitWebSocketEvents(this)
-                    }
+                val webSocketSession = authenticatedWebSocketClient.createWebSocketSession(clientId) {
+                    setWSSUrl(Url(serverLinks.webSocket), PATH_AWAIT)
+                    parameter(CLIENT_QUERY_KEY, clientId)
+                }
+                emitWebSocketEvents(webSocketSession)
             }
         }
 
     private suspend fun FlowCollector<WebSocketEvent<EventResponse>>.emitWebSocketEvents(
-        defaultClientWebSocketSession: DefaultClientWebSocketSession
+        defaultClientWebSocketSession: WebSocketSession
     ) {
         val logger = kaliumLogger.withFeatureId(EVENT_RECEIVER)
         logger.i("Websocket open")

network/src/commonMain/kotlin/com/wire/kalium/network/api/v0/authenticated/networkContainer/AuthenticatedNetworkContainerV0.kt

@@ -67,11 +67,13 @@ import com.wire.kalium.network.networkContainer.AuthenticatedNetworkContainer
 import com.wire.kalium.network.session.CertificatePinning
 import com.wire.kalium.network.session.SessionManager
 import io.ktor.client.engine.HttpClientEngine
+import io.ktor.websocket.WebSocketSession
 
 internal class AuthenticatedNetworkContainerV0 internal constructor(
     private val sessionManager: SessionManager,
     certificatePinning: CertificatePinning,
     mockEngine: HttpClientEngine?,
+    mockWebSocketSession: WebSocketSession?,
     kaliumLogger: KaliumLogger,
     engine: HttpClientEngine = mockEngine ?: defaultHttpEngine(
         serverConfigDTOApiProxy = sessionManager.serverConfig().links.apiProxy,
@@ -83,7 +85,12 @@ internal class AuthenticatedNetworkContainerV0 internal constructor(
         sessionManager = sessionManager,
         accessTokenApi = { httpClient -> AccessTokenApiV0(httpClient) },
         engine = engine,
-        kaliumLogger = kaliumLogger
+        kaliumLogger = kaliumLogger,
+        webSocketSessionProvider = if (mockWebSocketSession != null) {
+            { _, _ -> mockWebSocketSession }
+        } else {
+            null
+        }
     ) {
 
     override val accessTokenApi: AccessTokenApi get() = AccessTokenApiV0(networkClient.httpClient)

network/src/commonMain/kotlin/com/wire/kalium/network/api/v2/authenticated/networkContainer/AuthenticatedNetworkContainerV2.kt

@@ -68,13 +68,15 @@ import com.wire.kalium.network.networkContainer.AuthenticatedNetworkContainer
 import com.wire.kalium.network.session.CertificatePinning
 import com.wire.kalium.network.session.SessionManager
 import io.ktor.client.engine.HttpClientEngine
+import io.ktor.websocket.WebSocketSession
 
 @Suppress("LongParameterList")
 internal class AuthenticatedNetworkContainerV2 internal constructor(
     private val sessionManager: SessionManager,
     private val selfUserId: UserId,
     certificatePinning: CertificatePinning,
     mockEngine: HttpClientEngine?,
+    mockWebSocketSession: WebSocketSession?,
     kaliumLogger: KaliumLogger,
     engine: HttpClientEngine = mockEngine ?: defaultHttpEngine(
         serverConfigDTOApiProxy = sessionManager.serverConfig().links.apiProxy,
@@ -86,7 +88,12 @@ internal class AuthenticatedNetworkContainerV2 internal constructor(
         sessionManager = sessionManager,
         accessTokenApi = { httpClient -> AccessTokenApiV2(httpClient) },
         engine = engine,
-        kaliumLogger = kaliumLogger
+        kaliumLogger = kaliumLogger,
+        webSocketSessionProvider = if (mockWebSocketSession != null) {
+            { _, _ -> mockWebSocketSession }
+        } else {
+            null
+        }
     ) {
 
     override val accessTokenApi: AccessTokenApi get() = AccessTokenApiV2(networkClient.httpClient)

network/src/commonMain/kotlin/com/wire/kalium/network/api/v3/authenticated/networkContainer/AuthenticatedNetworkContainerV3.kt

@@ -69,13 +69,15 @@ import com.wire.kalium.network.networkContainer.AuthenticatedNetworkContainer
 import com.wire.kalium.network.session.CertificatePinning
 import com.wire.kalium.network.session.SessionManager
 import io.ktor.client.engine.HttpClientEngine
+import io.ktor.websocket.WebSocketSession
 
 @Suppress("LongParameterList")
 internal class AuthenticatedNetworkContainerV3 internal constructor(
     private val sessionManager: SessionManager,
     private val selfUserId: UserId,
     certificatePinning: CertificatePinning,
     mockEngine: HttpClientEngine?,
+    mockWebSocketSession: WebSocketSession?,
     kaliumLogger: KaliumLogger,
     engine: HttpClientEngine = mockEngine ?: defaultHttpEngine(
         serverConfigDTOApiProxy = sessionManager.serverConfig().links.apiProxy,
@@ -87,7 +89,12 @@ internal class AuthenticatedNetworkContainerV3 internal constructor(
         sessionManager = sessionManager,
         accessTokenApi = { httpClient -> AccessTokenApiV3(httpClient) },
         engine = engine,
-        kaliumLogger = kaliumLogger
+        kaliumLogger = kaliumLogger,
+        webSocketSessionProvider = if (mockWebSocketSession != null) {
+            { _, _ -> mockWebSocketSession }
+        } else {
+            null
+        }
     ) {
 
     override val accessTokenApi: AccessTokenApi get() = AccessTokenApiV3(networkClient.httpClient)

network/src/commonMain/kotlin/com/wire/kalium/network/api/v4/authenticated/networkContainer/AuthenticatedNetworkContainerV4.kt

@@ -68,13 +68,15 @@ import com.wire.kalium.network.networkContainer.AuthenticatedNetworkContainer
 import com.wire.kalium.network.session.CertificatePinning
 import com.wire.kalium.network.session.SessionManager
 import io.ktor.client.engine.HttpClientEngine
+import io.ktor.websocket.WebSocketSession
 
 @Suppress("LongParameterList")
 internal class AuthenticatedNetworkContainerV4 internal constructor(
     private val sessionManager: SessionManager,
     private val selfUserId: UserId,
     certificatePinning: CertificatePinning,
     mockEngine: HttpClientEngine?,
+    mockWebSocketSession: WebSocketSession?,
     kaliumLogger: KaliumLogger,
     engine: HttpClientEngine = mockEngine ?: defaultHttpEngine(
         serverConfigDTOApiProxy = sessionManager.serverConfig().links.apiProxy,
@@ -86,7 +88,12 @@ internal class AuthenticatedNetworkContainerV4 internal constructor(
         sessionManager = sessionManager,
         accessTokenApi = { httpClient -> AccessTokenApiV4(httpClient) },
         engine = engine,
-        kaliumLogger = kaliumLogger
+        kaliumLogger = kaliumLogger,
+        webSocketSessionProvider = if (mockWebSocketSession != null) {
+            { _, _ -> mockWebSocketSession }
+        } else {
+            null
+        }
     ) {
 
     override val accessTokenApi: AccessTokenApi get() = AccessTokenApiV4(networkClient.httpClient)

network/src/commonMain/kotlin/com/wire/kalium/network/api/v5/authenticated/networkContainer/AuthenticatedNetworkContainerV5.kt

@@ -68,13 +68,15 @@ import com.wire.kalium.network.networkContainer.AuthenticatedNetworkContainer
 import com.wire.kalium.network.session.CertificatePinning
 import com.wire.kalium.network.session.SessionManager
 import io.ktor.client.engine.HttpClientEngine
+import io.ktor.websocket.WebSocketSession
 
 @Suppress("LongParameterList")
 internal class AuthenticatedNetworkContainerV5 internal constructor(
     private val sessionManager: SessionManager,
     private val selfUserId: UserId,
     certificatePinning: CertificatePinning,
     mockEngine: HttpClientEngine?,
+    mockWebSocketSession: WebSocketSession?,
     kaliumLogger: KaliumLogger,
     engine: HttpClientEngine = mockEngine ?: defaultHttpEngine(
         serverConfigDTOApiProxy = sessionManager.serverConfig().links.apiProxy,
@@ -86,7 +88,12 @@ internal class AuthenticatedNetworkContainerV5 internal constructor(
         sessionManager = sessionManager,
         accessTokenApi = { httpClient -> AccessTokenApiV5(httpClient) },
         engine = engine,
-        kaliumLogger = kaliumLogger
+        kaliumLogger = kaliumLogger,
+        webSocketSessionProvider = if (mockWebSocketSession != null) {
+            { _, _ -> mockWebSocketSession }
+        } else {
+            null
+        }
     ) {
 
     override val accessTokenApi: AccessTokenApi get() = AccessTokenApiV5(networkClient.httpClient)

network/src/commonMain/kotlin/com/wire/kalium/network/api/v6/authenticated/networkContainer/AuthenticatedNetworkContainerV6.kt

@@ -68,13 +68,15 @@ import com.wire.kalium.network.networkContainer.AuthenticatedNetworkContainer
 import com.wire.kalium.network.session.CertificatePinning
 import com.wire.kalium.network.session.SessionManager
 import io.ktor.client.engine.HttpClientEngine
+import io.ktor.websocket.WebSocketSession
 
 @Suppress("LongParameterList")
 internal class AuthenticatedNetworkContainerV6 internal constructor(
     private val sessionManager: SessionManager,
     private val selfUserId: UserId,
     certificatePinning: CertificatePinning,
     mockEngine: HttpClientEngine?,
+    mockWebSocketSession: WebSocketSession?,
     kaliumLogger: KaliumLogger,
     engine: HttpClientEngine = mockEngine ?: defaultHttpEngine(
         serverConfigDTOApiProxy = sessionManager.serverConfig().links.apiProxy,
@@ -86,7 +88,12 @@ internal class AuthenticatedNetworkContainerV6 internal constructor(
         sessionManager = sessionManager,
         accessTokenApi = { httpClient -> AccessTokenApiV6(httpClient) },
         engine = engine,
-        kaliumLogger = kaliumLogger
+        kaliumLogger = kaliumLogger,
+        webSocketSessionProvider = if (mockWebSocketSession != null) {
+            { _, _ -> mockWebSocketSession }
+        } else {
+            null
+        }
     ) {
 
     override val accessTokenApi: AccessTokenApi get() = AccessTokenApiV6(networkClient.httpClient)