MlKitTestLauncher.kt

package com.google.mlkit.genai.prompt.test


import android.app.KeyguardManager
import android.graphics.Color
import android.graphics.Typeface
import android.view.Gravity
import android.view.WindowManager
import android.widget.TextView
import androidx.activity.ComponentActivity
import androidx.core.content.getSystemService
import androidx.lifecycle.Lifecycle
import androidx.test.core.app.ActivityScenario
import com.google.mlkit.genai.common.DownloadStatus
import com.google.mlkit.genai.common.FeatureStatus
import com.google.mlkit.genai.prompt.Generation
import com.google.mlkit.genai.prompt.GenerationConfig
import com.google.mlkit.genai.prompt.GenerativeModel
import kotlinx.coroutines.test.runTest
import kotlinx.coroutines.withTimeout
import kotlinx.coroutines.yield
import kotlin.coroutines.CoroutineContext
import kotlin.coroutines.EmptyCoroutineContext
import kotlin.time.Duration
import kotlin.time.Duration.Companion.minutes
import kotlin.time.Duration.Companion.seconds


/**
* A custom test runner wrapper designed for integration tests that utilize ML Kit's GenAI APIs (Gemini Nano).
*
* @param generationConfig Optional [GenerationConfig] to configure the model instance (e.g. Setting temperature, topK, etc.)
* @param context The CoroutineContext for the test execution.
* @param timeout The maximum allowed [Duration] for the test to run. Defaults to 5 minutes because downloading the model can take considerable time.
* @param testBody The suspension block containing the test logic, executed securely with a ready [GenerativeModel].
*/
fun runPromptApiTest(
   generationConfig: GenerationConfig? = null,
   context: CoroutineContext = EmptyCoroutineContext,
   timeoutWithDownload: Duration = 5.minutes, // Only once per device
   timeout: Duration = 30.seconds,
   testBody: suspend (GenerativeModel) -> Unit,
) = runForegroundTest(context = context, timeout = timeoutWithDownload) {
   val model = if (generationConfig != null) {
       Generation.getClient(generationConfig)
   } else {
       Generation.getClient()
   }
   try {
       model.ensureReady()
       withTimeout(timeout) {
           testBody(model)
       }
   } finally {
       model.close()
   }
}


/**
* Ensures that the [GenerativeModel] is fully downloaded and available on the device before continuing execution.
*
* If the model is not available, it triggers a download and suspends the coroutine until the download stream
* completes successfully. If the device does not support Gemini Nano, it throws an [IllegalStateException].
*/
suspend fun GenerativeModel.ensureReady() {
   val status = checkStatus()


   when (status) {
       FeatureStatus.AVAILABLE -> return // Already good to go
       FeatureStatus.UNAVAILABLE -> error("Gemini Nano not supported on this device.")
       FeatureStatus.DOWNLOADING, FeatureStatus.DOWNLOADABLE -> {
           // Trigger download and suspend until the Flow completes/fails
           download().collect { downloadStatus ->
               when (downloadStatus) {
                   is DownloadStatus.DownloadCompleted -> return@collect
                   is DownloadStatus.DownloadFailed -> throw downloadStatus.e
                   else -> Unit // Ignore progress updates
               }
           }
       }
   }
}


/**
* A custom test runner wrapper that ensures test execution happens while the app is in the foreground.
*
* Some Android libraries and system services (like ML Kit's GenAI Prompt API / AICore) strictly enforce
* foreground usage policies and will throw exceptions in headless tests:
* ```
* Background usage is blocked. Please use the API when your app is in the foreground instead.
* ```
*
* @param context The CoroutineContext for the test execution.
* @param timeout The maximum allowed [Duration] for the test to run. Defaults to 30 seconds (overriding the default 10 seconds of runTest).
* @param testBody The suspension block containing the test logic, executed securely in the foreground.
*/
fun runForegroundTest(
   context: CoroutineContext = EmptyCoroutineContext,
   timeout: Duration = 30.seconds,
   testBody: suspend () -> Unit,
) = runTest(context = context, timeout = timeout) {
   ActivityScenario.launch(ComponentActivity::class.java).use { scenario ->


       scenario.onActivity { activity ->
           with(activity) {
               // 1. Force the screen on and unlock
               setShowWhenLocked(true)
               setTurnScreenOn(true)
               window.addFlags(WindowManager.LayoutParams.FLAG_KEEP_SCREEN_ON)
               window.decorView.setBackgroundColor(Color.BLACK)
               setContentView(
                   TextView(this).apply {
                       setTextColor(Color.GREEN)
                       setBackgroundColor(Color.BLACK)
                       textSize = 16f
                       typeface = Typeface.MONOSPACE
                       gravity = Gravity.CENTER
                       text = "Running Test in Foreground..."
                   },
               )
               getSystemService<KeyguardManager>()?.requestDismissKeyguard(activity, null)
           }
       }


       // 2. Wait for Resume
       scenario.moveToState(Lifecycle.State.RESUMED)


       // 3. Brief delay to ensure Window Focus is gained
       yield()


       // 4. Execute the test logic in the foreground
       testBody()
   }
}