Android test use kotlin

extension/android/executorch_android/build.gradle

@@ -9,6 +9,7 @@
 plugins {
     id "com.android.library" version "8.9.0"
     id "com.vanniktech.maven.publish" version "0.31.0"
+    alias(libs.plugins.jetbrains.kotlin.android)
 }
 
 android {
@@ -34,6 +35,9 @@ android {
             resources.srcDirs += [ 'src/androidTest/resources' ]
         }
     }
+    kotlinOptions {
+        jvmTarget = '1.8'
+    }
 }
 
 task copyTestRes(type: Exec) {
@@ -43,6 +47,7 @@ task copyTestRes(type: Exec) {
 dependencies {
     implementation 'com.facebook.fbjni:fbjni:0.5.1'
     implementation 'com.facebook.soloader:nativeloader:0.10.5'
+    implementation libs.core.ktx
     testImplementation 'junit:junit:4.12'
     androidTestImplementation 'androidx.test.ext:junit:1.1.5'
     androidTestImplementation 'androidx.test:rules:1.2.0'

extension/android/executorch_android/src/androidTest/java/org/pytorch/executorch/LlmModuleInstrumentationTest.kt

@@ -0,0 +1,115 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+package org.pytorch.executorch
+
+import android.Manifest
+import androidx.test.InstrumentationRegistry
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.rule.GrantPermissionRule
+import org.apache.commons.io.FileUtils
+import org.json.JSONException
+import org.json.JSONObject
+import org.junit.Assert
+import org.junit.Before
+import org.junit.Rule
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.pytorch.executorch.extension.llm.LlmCallback
+import org.pytorch.executorch.extension.llm.LlmModule
+import java.io.File
+import java.io.IOException
+import java.net.URISyntaxException
+
+/** Unit tests for [org.pytorch.executorch.extension.llm.LlmModule].  */
+@RunWith(AndroidJUnit4::class)
+class LlmModuleInstrumentationTest : LlmCallback {
+    private val results: MutableList<String> = ArrayList()
+    private val tokensPerSecond: MutableList<Float> = ArrayList()
+    private var mModule: LlmModule? = null
+
+    @Before
+    @Throws(IOException::class)
+    fun setUp() {
+        // copy zipped test resources to local device
+        val addPteFile = File(getTestFilePath(TEST_FILE_NAME))
+        var inputStream = javaClass.getResourceAsStream(TEST_FILE_NAME)
+        FileUtils.copyInputStreamToFile(inputStream, addPteFile)
+        inputStream.close()
+
+        val tokenizerFile = File(getTestFilePath(TOKENIZER_FILE_NAME))
+        inputStream = javaClass.getResourceAsStream(TOKENIZER_FILE_NAME)
+        FileUtils.copyInputStreamToFile(inputStream, tokenizerFile)
+        inputStream.close()
+
+        mModule =
+            LlmModule(getTestFilePath(TEST_FILE_NAME), getTestFilePath(TOKENIZER_FILE_NAME), 0.0f)
+    }
+
+    @Rule
+    var mRuntimePermissionRule: GrantPermissionRule =
+        GrantPermissionRule.grant(Manifest.permission.READ_EXTERNAL_STORAGE)
+
+    @Test
+    @Throws(IOException::class, URISyntaxException::class)
+    fun testGenerate() {
+        val loadResult = mModule!!.load()
+        // Check that the model can be load successfully
+        Assert.assertEquals(OK.toLong(), loadResult.toLong())
+
+        mModule!!.generate(TEST_PROMPT, SEQ_LEN, this@LlmModuleInstrumentationTest)
+        Assert.assertEquals(results.size.toLong(), SEQ_LEN.toLong())
+        Assert.assertTrue(tokensPerSecond[tokensPerSecond.size - 1] > 0)
+    }
+
+    @Test
+    @Throws(IOException::class, URISyntaxException::class)
+    fun testGenerateAndStop() {
+        mModule!!.generate(TEST_PROMPT, SEQ_LEN, object : LlmCallback {
+            override fun onResult(result: String) {
+                [email protected](result)
+                mModule!!.stop()
+            }
+
+            override fun onStats(stats: String) {
+                [email protected](stats)
+            }
+        })
+
+        val stoppedResultSize = results.size
+        Assert.assertTrue(stoppedResultSize < SEQ_LEN)
+    }
+
+    override fun onResult(result: String) {
+        results.add(result)
+    }
+
+    override fun onStats(stats: String) {
+        var tps = 0f
+        try {
+            val jsonObject = JSONObject(stats)
+            val numGeneratedTokens = jsonObject.getInt("generated_tokens")
+            val inferenceEndMs = jsonObject.getInt("inference_end_ms")
+            val promptEvalEndMs = jsonObject.getInt("prompt_eval_end_ms")
+            tps = numGeneratedTokens.toFloat() / (inferenceEndMs - promptEvalEndMs) * 1000
+            tokensPerSecond.add(tps)
+        } catch (e: JSONException) {
+        }
+    }
+
+    companion object {
+        private const val TEST_FILE_NAME = "/stories.pte"
+        private const val TOKENIZER_FILE_NAME = "/tokenizer.bin"
+        private const val TEST_PROMPT = "Hello"
+        private const val OK = 0x00
+        private const val SEQ_LEN = 32
+
+        private fun getTestFilePath(fileName: String): String {
+            return InstrumentationRegistry.getInstrumentation().targetContext.externalCacheDir.toString() + fileName
+        }
+    }
+}