From 97d95758beed33948a928823482cbfff5dab73f5 Mon Sep 17 00:00:00 2001
From: David Motsonashvili <davidmotson@google.com>
Date: Tue, 23 Jul 2024 09:33:35 -0700
Subject: [PATCH] change the test constructor to allow injecting an http engine
 to allow more configuration of the mockengine

---
 common/build.gradle.kts                                |  2 +-
 .../ai/client/generativeai/common/APIController.kt     | 10 ++--------
 .../google/ai/client/generativeai/common/util/tests.kt |  7 +++++--
 3 files changed, 8 insertions(+), 11 deletions(-)

diff --git a/common/build.gradle.kts b/common/build.gradle.kts
index 842e9301..406b7eee 100644
--- a/common/build.gradle.kts
+++ b/common/build.gradle.kts
@@ -90,7 +90,7 @@ dependencies {
     implementation("com.google.guava:listenablefuture:1.0")
     implementation("androidx.concurrent:concurrent-futures:1.2.0-alpha02")
     implementation("androidx.concurrent:concurrent-futures-ktx:1.2.0-alpha02")
-    compileOnly("io.ktor:ktor-client-mock:$ktorVersion")
+
     testImplementation("junit:junit:4.13.2")
     testImplementation("io.kotest:kotest-assertions-core:4.0.7")
     testImplementation("io.kotest:kotest-assertions-jvm:4.0.7")
diff --git a/common/src/main/kotlin/com/google/ai/client/generativeai/common/APIController.kt b/common/src/main/kotlin/com/google/ai/client/generativeai/common/APIController.kt
index aaf79634..011e1d4c 100644
--- a/common/src/main/kotlin/com/google/ai/client/generativeai/common/APIController.kt
+++ b/common/src/main/kotlin/com/google/ai/client/generativeai/common/APIController.kt
@@ -24,8 +24,6 @@ import com.google.ai.client.generativeai.common.util.fullModelName
 import io.ktor.client.HttpClient
 import io.ktor.client.call.body
 import io.ktor.client.engine.HttpClientEngine
-import io.ktor.client.engine.mock.MockEngine
-import io.ktor.client.engine.mock.respond
 import io.ktor.client.engine.okhttp.OkHttp
 import io.ktor.client.plugins.HttpTimeout
 import io.ktor.client.plugins.contentnegotiation.ContentNegotiation
@@ -38,12 +36,9 @@ import io.ktor.client.statement.HttpResponse
 import io.ktor.client.statement.bodyAsChannel
 import io.ktor.client.statement.bodyAsText
 import io.ktor.http.ContentType
-import io.ktor.http.HttpHeaders
 import io.ktor.http.HttpStatusCode
 import io.ktor.http.contentType
-import io.ktor.http.headersOf
 import io.ktor.serialization.kotlinx.json.json
-import io.ktor.utils.io.ByteChannel
 import kotlin.time.Duration
 import kotlinx.coroutines.CoroutineName
 import kotlinx.coroutines.TimeoutCancellationException
@@ -98,13 +93,12 @@ internal constructor(
     requestOptions: RequestOptions,
     apiClient: String,
     headerProvider: HeaderProvider?,
-    channel: ByteChannel,
-    status: HttpStatusCode,
+    engine: HttpClientEngine,
   ) : this(
     key,
     model,
     requestOptions,
-    MockEngine { respond(channel, status, headersOf(HttpHeaders.ContentType, "application/json")) },
+    engine,
     apiClient,
     headerProvider,
   )
diff --git a/common/src/test/java/com/google/ai/client/generativeai/common/util/tests.kt b/common/src/test/java/com/google/ai/client/generativeai/common/util/tests.kt
index 0726d959..d0e989be 100644
--- a/common/src/test/java/com/google/ai/client/generativeai/common/util/tests.kt
+++ b/common/src/test/java/com/google/ai/client/generativeai/common/util/tests.kt
@@ -28,7 +28,11 @@ import com.google.ai.client.generativeai.common.shared.Content
 import com.google.ai.client.generativeai.common.shared.TextPart
 import io.kotest.matchers.collections.shouldNotBeEmpty
 import io.kotest.matchers.nulls.shouldNotBeNull
+import io.ktor.client.engine.mock.MockEngine
+import io.ktor.client.engine.mock.respond
+import io.ktor.http.HttpHeaders
 import io.ktor.http.HttpStatusCode
+import io.ktor.http.headersOf
 import io.ktor.utils.io.ByteChannel
 import io.ktor.utils.io.close
 import io.ktor.utils.io.writeFully
@@ -108,8 +112,7 @@ internal fun commonTest(
       requestOptions,
       TEST_CLIENT_ID,
       null,
-      channel,
-      status,
+      MockEngine { respond(channel, status, headersOf(HttpHeaders.ContentType, "application/json")) },
     )
   CommonTestScope(channel, apiController).block()
 }