minor restructure and added tests

Author: David Motsonashvili

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

@@ -16,6 +16,7 @@
 
 package com.google.firebase.vertexai.common
 
+import android.content.Context
 import android.util.Log
 import com.google.firebase.Firebase
 import com.google.firebase.FirebaseApp
@@ -93,6 +94,8 @@ internal constructor(
   httpEngine: HttpClientEngine,
   private val apiClient: String,
   private val firebaseApp: FirebaseApp,
+  private val appVersion: Int = 0,
+  private val googleAppId: String,
   private val headerProvider: HeaderProvider?,
 ) {
 
@@ -103,7 +106,17 @@ internal constructor(
     apiClient: String,
     firebaseApp: FirebaseApp,
     headerProvider: HeaderProvider? = null,
-  ) : this(key, model, requestOptions, OkHttp.create(), apiClient, firebaseApp, headerProvider)
+  ) : this(
+    key,
+    model,
+    requestOptions,
+    OkHttp.create(),
+    apiClient,
+    firebaseApp,
+    getVersionNumber(firebaseApp),
+    firebaseApp.options.applicationId,
+    headerProvider
+  )
 
   private val model = fullModelName(model)
 
@@ -179,20 +192,11 @@ internal constructor(
     header("x-goog-api-key", key)
     header("x-goog-api-client", apiClient)
     if (firebaseApp.isDataCollectionDefaultEnabled) {
-      header("X-Firebase-AppId", firebaseApp.options.applicationId)
-      header("X-Firebase-AppVersion", getVersionNumber())
+      header("X-Firebase-AppId", googleAppId)
+      header("X-Firebase-AppVersion", appVersion)
     }
   }
 
-  private fun getVersionNumber(): Int {
-    try {
-      val context = firebaseApp.applicationContext
-      return context.packageManager.getPackageInfo(context.packageName, 0).versionCode
-    } catch (e: Exception) {
-      Log.d(TAG, "Error while getting app version: ${e.message}")
-      return 0
-    }
-  }
 
   private suspend fun HttpRequestBuilder.applyHeaderProvider() {
     if (headerProvider != null) {
@@ -257,6 +261,16 @@ internal constructor(
 
   companion object {
     private val TAG = APIController::class.java.simpleName
+
+    private fun getVersionNumber(app: FirebaseApp): Int {
+      try {
+        val context = app.applicationContext
+        return context.packageManager.getPackageInfo(context.packageName, 0).versionCode
+      } catch (e: Exception) {
+        Log.d(TAG, "Error while getting app version: ${e.message}")
+        return 0
+      }
+    }
   }
 }
 

firebase-vertexai/src/test/java/com/google/firebase/vertexai/GenerativeModelTesting.kt

@@ -49,6 +49,8 @@ import org.mockito.Mockito
 
 internal class GenerativeModelTesting {
   private val TEST_CLIENT_ID = "test"
+  private val TEST_APP_ID = "1:android:12345"
+  private val TEST_VERSION = 1
 
   private var mockFirebaseApp: FirebaseApp = Mockito.mock<FirebaseApp>()
 
@@ -75,6 +77,8 @@ internal class GenerativeModelTesting {
         mockEngine,
         TEST_CLIENT_ID,
         mockFirebaseApp,
+        TEST_VERSION,
+        TEST_APP_ID,
         null,
       )
 
@@ -121,6 +125,8 @@ internal class GenerativeModelTesting {
         mockEngine,
         TEST_CLIENT_ID,
         mockFirebaseApp,
+        TEST_VERSION,
+        TEST_APP_ID,
         null,
       )
 

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

@@ -16,7 +16,11 @@
 
 package com.google.firebase.vertexai.common
 
+import android.content.Context
+import android.content.pm.PackageInfo
+import android.content.pm.PackageManager
 import com.google.firebase.FirebaseApp
+import com.google.firebase.FirebaseOptions
 import com.google.firebase.vertexai.BuildConfig
 import com.google.firebase.vertexai.common.util.commonTest
 import com.google.firebase.vertexai.common.util.createResponses
@@ -54,10 +58,16 @@ import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
 import org.junit.runners.Parameterized
+import org.mockito.ArgumentMatchers.any
+import org.mockito.ArgumentMatchers.anyInt
 import org.mockito.Mockito
 
 private val TEST_CLIENT_ID = "genai-android/test"
 
+private val TEST_APP_ID = "1:android:12345"
+
+private val TEST_VERSION = 1
+
 internal class APIControllerTests {
   private val testTimeout = 5.seconds
 
@@ -113,6 +123,8 @@ internal class RequestFormatTests {
         mockEngine,
         "genai-android/${BuildConfig.VERSION_NAME}",
         mockFirebaseApp,
+        TEST_VERSION,
+        TEST_APP_ID,
         null,
       )
 
@@ -141,6 +153,8 @@ internal class RequestFormatTests {
         mockEngine,
         TEST_CLIENT_ID,
         mockFirebaseApp,
+        TEST_VERSION,
+        TEST_APP_ID,
         null,
       )
 
@@ -169,6 +183,8 @@ internal class RequestFormatTests {
         mockEngine,
         TEST_CLIENT_ID,
         mockFirebaseApp,
+        TEST_VERSION,
+        TEST_APP_ID,
         null,
       )
 
@@ -177,6 +193,35 @@ internal class RequestFormatTests {
     mockEngine.requestHistory.first().headers["x-goog-api-client"] shouldBe TEST_CLIENT_ID
   }
 
+  @Test
+  fun `ml monitoring header is set correctly if data collection is enabled`() = doBlocking {
+    val response = JSON.encodeToString(CountTokensResponse.Internal(totalTokens = 10))
+    val mockEngine = MockEngine {
+      respond(response, HttpStatusCode.OK, headersOf(HttpHeaders.ContentType, "application/json"))
+    }
+
+    Mockito.`when`(mockFirebaseApp.isDataCollectionDefaultEnabled).thenReturn(true)
+
+    val controller =
+      APIController(
+        "super_cool_test_key",
+        "gemini-pro-1.5",
+        RequestOptions(),
+        mockEngine,
+        TEST_CLIENT_ID,
+        mockFirebaseApp,
+        TEST_VERSION,
+        TEST_APP_ID,
+        null,
+      )
+
+    withTimeout(5.seconds) { controller.countTokens(textCountTokenRequest("cats")) }
+
+    mockEngine.requestHistory.first().headers["X-Firebase-AppId"] shouldBe TEST_APP_ID
+    mockEngine.requestHistory.first().headers["X-Firebase-AppVersion"] shouldBe
+      TEST_VERSION.toString()
+  }
+
   @Test
   fun `ToolConfig serialization contains correct keys`() = doBlocking {
     val channel = ByteChannel(autoFlush = true)
@@ -193,6 +238,8 @@ internal class RequestFormatTests {
         mockEngine,
         TEST_CLIENT_ID,
         mockFirebaseApp,
+        TEST_VERSION,
+        TEST_APP_ID,
         null,
       )
 
@@ -245,6 +292,8 @@ internal class RequestFormatTests {
         mockEngine,
         TEST_CLIENT_ID,
         mockFirebaseApp,
+        TEST_VERSION,
+        TEST_APP_ID,
         testHeaderProvider,
       )
 
@@ -280,6 +329,8 @@ internal class RequestFormatTests {
         mockEngine,
         TEST_CLIENT_ID,
         mockFirebaseApp,
+        TEST_VERSION,
+        TEST_APP_ID,
         testHeaderProvider,
       )
 
@@ -304,6 +355,8 @@ internal class RequestFormatTests {
         mockEngine,
         TEST_CLIENT_ID,
         mockFirebaseApp,
+        TEST_VERSION,
+        TEST_APP_ID,
         null,
       )
 
@@ -349,6 +402,8 @@ internal class ModelNamingTests(private val modelName: String, private val actua
         mockEngine,
         TEST_CLIENT_ID,
         mockFirebaseApp,
+        TEST_VERSION,
+        TEST_APP_ID,
         null,
       )
 

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

@@ -37,6 +37,8 @@ import kotlinx.serialization.encodeToString
 import org.mockito.Mockito
 
 private val TEST_CLIENT_ID = "genai-android/test"
+private val TEST_APP_ID = "1:android:12345"
+private val TEST_VERSION = 1
 
 internal fun prepareStreamingResponse(
   response: List<GenerateContentResponse.Internal>
@@ -106,6 +108,8 @@ internal fun commonTest(
       },
       TEST_CLIENT_ID,
       mockFirebaseApp,
+      TEST_VERSION,
+      TEST_APP_ID,
       null,
     )
   CommonTestScope(channel, apiController).block()

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

@@ -39,6 +39,8 @@ import kotlinx.coroutines.launch
 import org.mockito.Mockito
 
 private val TEST_CLIENT_ID = "firebase-vertexai-android/test"
+private val TEST_APP_ID = "1:android:12345"
+private val TEST_VERSION = 1
 
 /** String separator used in SSE communication to signal the end of a message. */
 internal const val SSE_SEPARATOR = "\r\n\r\n"
@@ -115,6 +117,8 @@ internal fun commonTest(
       },
       TEST_CLIENT_ID,
       mockFirebaseApp,
+      TEST_VERSION,
+      TEST_APP_ID,
       null,
     )
   val model = GenerativeModel("cool-model-name", controller = apiController)