Refactor TokenPlugin to use createClientPlugin

network/src/commonMain/kotlin/dev/icerock/moko/network/plugins/TokenPlugin.kt

@@ -4,40 +4,35 @@
 
 package dev.icerock.moko.network.plugins
 
-import io.ktor.client.HttpClient
-import io.ktor.client.plugins.HttpClientPlugin
-import io.ktor.client.request.HttpRequestPipeline
+import io.ktor.client.plugins.createClientPlugin
 import io.ktor.client.request.header
-import io.ktor.util.AttributeKey
 
-class TokenPlugin private constructor(
-    private val tokenHeaderName: String,
-    private val tokenProvider: TokenProvider
+val TokenPlugin = createClientPlugin(
+    name = "TokenPlugin",
+    createConfiguration = ::TokenPluginConfig
 ) {
+    val tokenHeaderName = pluginConfig.tokenHeaderName
+    val tokenProvider = pluginConfig.tokenProvider
 
-    class Config {
-        var tokenHeaderName: String? = null
-        var tokenProvider: TokenProvider? = null
-        fun build() = TokenPlugin(
-            tokenHeaderName ?: throw IllegalArgumentException("HeaderName should be contain"),
-            tokenProvider ?: throw IllegalArgumentException("TokenProvider should be contain")
-        )
+    if (tokenHeaderName == null) {
+        throw IllegalArgumentException("HeaderName should be contain")
     }
 
-    companion object Plugin : HttpClientPlugin<Config, TokenPlugin> {
-        override val key = AttributeKey<TokenPlugin>("TokenPlugin")
-
-        override fun prepare(block: Config.() -> Unit) = Config().apply(block).build()
+    if (tokenProvider == null) {
+        throw IllegalArgumentException("TokenProvider should be contain")
+    }
 
-        override fun install(plugin: TokenPlugin, scope: HttpClient) {
-            scope.requestPipeline.intercept(HttpRequestPipeline.State) {
-                plugin.tokenProvider.getToken()?.apply {
-                    context.headers.remove(plugin.tokenHeaderName)
-                    context.header(plugin.tokenHeaderName, this)
-                }
-            }
+    onRequest { request, _ ->
+        tokenProvider.getToken()?.apply {
+            request.headers.remove(tokenHeaderName)
+            request.headers.append(tokenHeaderName, this)
         }
     }
+}
+
+class TokenPluginConfig {
+    var tokenHeaderName: String? = null
+    var tokenProvider: TokenProvider? = null
 
     fun interface TokenProvider {
         fun getToken(): String?

network/src/commonTest/kotlin/dev/icerock/moko/network/plugins/TokenPluginTest.kt

@@ -0,0 +1,85 @@
+/*
+ * Copyright 2024 IceRock MAG Inc. Use of this source code is governed by the Apache 2.0 license.
+ */
+
+package dev.icerock.moko.network.plugins
+
+import dev.icerock.moko.network.plugins.TokenPluginConfig.TokenProvider
+import io.ktor.client.HttpClient
+import io.ktor.client.engine.mock.MockEngine
+import io.ktor.client.engine.mock.respond
+import io.ktor.client.request.get
+import io.ktor.http.HttpHeaders
+import kotlinx.coroutines.test.runTest
+import kotlin.test.Test
+import kotlin.test.assertEquals
+import kotlin.test.assertFailsWith
+import kotlin.test.assertFalse
+import kotlin.test.assertNotNull
+import kotlin.test.assertTrue
+
+class TokenPluginTest {
+
+    private val mockToken = "mockTokenValue"
+    private val mockHeaderName = "X-Auth-Token"
+
+    @Test
+    fun `token successfully added`() = runTest {
+        val mockEngine = MockEngine { request ->
+            assertTrue(request.headers.contains(mockHeaderName, mockToken))
+            respond("")
+        }
+
+        val client = HttpClient(mockEngine) {
+            install(TokenPlugin) {
+                this.tokenHeaderName = mockHeaderName
+                this.tokenProvider = TokenProvider { mockToken }
+            }
+        }
+
+        client.get("http://localhost/test")
+    }
+
+    @Test
+    fun `null token`() = runTest {
+        val mockEngine = MockEngine { request ->
+            assertFalse(request.headers.contains(mockHeaderName))
+            respond("")
+        }
+
+        val client = HttpClient(mockEngine) {
+            install(TokenPlugin) {
+                this.tokenHeaderName = mockHeaderName
+                this.tokenProvider = TokenProvider { null }
+            }
+        }
+
+        client.get("http://localhost/test")
+    }
+
+    @Test
+    fun `config validation - null tokenHeaderName`() = runTest {
+        val exception = assertFailsWith<IllegalArgumentException> {
+            HttpClient(MockEngine { respond("") }) {
+                install(TokenPlugin) {
+                    this.tokenHeaderName = null
+                    this.tokenProvider = TokenProvider { mockToken }
+                }
+            }
+        }
+        assertEquals("HeaderName should be contain", exception.message)
+    }
+
+    @Test
+    fun `config validation - null tokenProvider`() = runTest {
+        val exception = assertFailsWith<IllegalArgumentException> {
+            HttpClient(MockEngine { respond("") }) {
+                install(TokenPlugin) {
+                    this.tokenHeaderName = mockHeaderName
+                    this.tokenProvider = null
+                }
+            }
+        }
+        assertEquals("TokenProvider should be contain", exception.message)
+    }
+}