Refactor out ktor mock to simplify rules

Author: emilypgoogle

firebase-vertexai/consumer-rules.pro

@@ -21,8 +21,3 @@
 #-renamesourcefileattribute SourceFile
 
 -keep class com.google.firebase.vertexai.common.** { *; }
-
--dontwarn io.ktor.client.engine.mock.MockEngine$Companion
--dontwarn io.ktor.client.engine.mock.MockEngine
--dontwarn io.ktor.client.engine.mock.MockRequestHandleScope
--dontwarn io.ktor.client.engine.mock.MockUtilsKt

firebase-vertexai/gradle.properties

@@ -12,5 +12,5 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-version=16.0.0-beta06
+version=16.0.0-beta10
 latestReleasedVersion=16.0.0-beta05

firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/common/APIController.kt

@@ -17,15 +17,12 @@
 package com.google.firebase.vertexai.common
 
 import android.util.Log
-import androidx.annotation.VisibleForTesting
 import com.google.firebase.vertexai.common.server.FinishReason
 import com.google.firebase.vertexai.common.util.decodeToFlow
 import com.google.firebase.vertexai.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,13 +35,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
 import kotlinx.coroutines.flow.Flow
@@ -54,6 +47,7 @@ import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.launch
 import kotlinx.coroutines.withTimeout
 import kotlinx.serialization.json.Json
+import kotlin.time.Duration
 
 internal val JSON = Json {
   ignoreUnknownKeys = true
@@ -91,24 +85,6 @@ internal constructor(
     headerProvider: HeaderProvider? = null,
   ) : this(key, model, requestOptions, OkHttp.create(), apiClient, headerProvider)
 
-  @VisibleForTesting(otherwise = VisibleForTesting.NONE)
-  constructor(
-    key: String,
-    model: String,
-    requestOptions: RequestOptions,
-    apiClient: String,
-    headerProvider: HeaderProvider?,
-    channel: ByteChannel,
-    status: HttpStatusCode,
-  ) : this(
-    key,
-    model,
-    requestOptions,
-    MockEngine { respond(channel, status, headersOf(HttpHeaders.ContentType, "application/json")) },
-    apiClient,
-    headerProvider,
-  )
-
   private val model = fullModelName(model)
 
   private val client =

firebase-vertexai/src/test/java/com/google/firebase/vertexai/common/util/tests.kt

@@ -28,7 +28,11 @@ import com.google.firebase.vertexai.common.shared.Content
 import com.google.firebase.vertexai.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
@@ -106,10 +110,9 @@ internal fun commonTest(
       "super_cool_test_key",
       "gemini-pro",
       requestOptions,
+      MockEngine { respond(channel, status, headersOf(HttpHeaders.ContentType, "application/json")) },
       TEST_CLIENT_ID,
       null,
-      channel,
-      status,
     )
   CommonTestScope(channel, apiController).block()
 }

firebase-vertexai/src/test/java/com/google/firebase/vertexai/util/tests.kt

@@ -21,7 +21,11 @@ import com.google.firebase.vertexai.common.APIController
 import com.google.firebase.vertexai.common.RequestOptions
 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
@@ -93,10 +97,9 @@ internal fun commonTest(
       "super_cool_test_key",
       "gemini-pro",
       requestOptions,
+      MockEngine { respond(channel, status, headersOf(HttpHeaders.ContentType, "application/json")) },
       TEST_CLIENT_ID,
       null,
-      channel,
-      status,
     )
   val model = GenerativeModel("cool-model-name", controller = apiController)
   CommonTestScope(channel, model).block()