LlamaDemo Android UI workflow automated test (#143)

Author: kirklandsign

llm/android/LlamaDemo/README.md

@@ -189,32 +189,55 @@ Ensure you have the following functions in your callback class that you provided
 
 ```
 
-## Instrumentation Test
-You can run the instrumentation test for a sanity check. The test loads a model .pte file and tokenizer.bin file
-under `/data/local/tmp/llama`.
+## Instrumentation Tests
+
+The app includes instrumentation tests for sanity checking and UI workflow validation.
+
+### Available Tests
+
+1. **SanityCheck** - Basic model loading and generation test that verifies the LLM module can load a model and generate tokens.
+
+2. **UIWorkflowTest** - UI-based tests that simulate user interactions:
+   - `testModelLoadingWorkflow`: Tests the complete flow of selecting a model/tokenizer and loading it
+   - `testSendMessageAndReceiveResponse`: Tests sending a message and receiving a response from the model
+
+### Model Preparation
+
+The test model (`stories110M.pte`) and tokenizer (`tokenizer.model`) are **automatically downloaded** when you run the tests via Gradle. The download task runs before the instrumentation tests execute.
+
+If you want to manually prepare the model files, you can use the following commands:
 
-### Model preparation
-You need to install [executorch python package](https://docs.pytorch.org/executorch/stable/getting-started.html#installation) first.
 ```sh
+# Install executorch python package first: https://docs.pytorch.org/executorch/stable/getting-started.html#installation
+
 curl -C - -Ls "https://huggingface.co/karpathy/tinyllamas/resolve/main/stories110M.pt" --output stories110M.pt
 curl -C - -Ls "https://raw.githubusercontent.com/karpathy/llama2.c/master/tokenizer.model" --output tokenizer.model
+
 # Create params.json file
 touch params.json
 echo '{"dim": 768, "multiple_of": 32, "n_heads": 12, "n_layers": 12, "norm_eps": 1e-05, "vocab_size": 32000}' > params.json
-python -m executorch.extension.llm.export.export_llm base.checkpoint=stories110M.pt base.params=params.json model.dtype_override="fp16" export.output_name=stories110m_h.pte model.use_kv_cache=True
-python -m pytorch_tokenizers.tools.llama2c.convert -t tokenizer.model -o tokenizer.bin
-```
-### Push model
-```sh
+
+# Export the model
+python -m executorch.extension.llm.export.export_llm base.checkpoint=stories110M.pt base.params=params.json model.dtype_override="fp16" export.output_name=stories110M.pte model.use_kv_cache=True
+
+# Push to device
 adb shell mkdir -p /data/local/tmp/llama
-adb push stories110m_h.pte /data/local/tmp/llama
-adb push tokenizer.bin /data/local/tmp/llama
+adb push stories110M.pte /data/local/tmp/llama
+adb push tokenizer.model /data/local/tmp/llama
 ```
 
-### Run test
+### Running Tests
+
+Run all instrumentation tests:
 ```sh
 ./gradlew connectedAndroidTest
 ```
 
+Run a specific test class:
+```sh
+./gradlew connectedAndroidTest -Pandroid.testInstrumentationRunnerArguments.class=com.example.executorchllamademo.SanityCheck
+./gradlew connectedAndroidTest -Pandroid.testInstrumentationRunnerArguments.class=com.example.executorchllamademo.UIWorkflowTest
+```
+
 ## Reporting Issues
 If you encountered any bugs or issues following this tutorial, please file a bug/issue here on [GitHub](https://github.com/pytorch/executorch/issues/new), or join our Discord [here](https://lnkd.in/gWCM4ViK).

llm/android/LlamaDemo/app/build.gradle.kts

@@ -196,6 +196,7 @@ dependencies {
   testImplementation("junit:junit:4.13.2")
   androidTestImplementation("androidx.test.ext:junit:1.1.5")
   androidTestImplementation("androidx.test.espresso:espresso-core:3.5.1")
+  androidTestImplementation("androidx.test.uiautomator:uiautomator:2.2.0")
   androidTestImplementation(platform("androidx.compose:compose-bom:2023.03.00"))
   androidTestImplementation("androidx.compose.ui:ui-test-junit4")
   debugImplementation("androidx.compose.ui:ui-tooling")

llm/android/LlamaDemo/app/src/androidTest/java/com/example/executorchllamademo/PerfTest.java

@@ -0,0 +1,58 @@
+/*
+ * 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 com.example.executorchllamademo;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+import java.io.File;
+import java.util.ArrayList;
+import java.util.List;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.pytorch.executorch.extension.llm.LlmCallback;
+import org.pytorch.executorch.extension.llm.LlmModule;
+
+@RunWith(AndroidJUnit4.class)
+public class SanityCheck implements LlmCallback {
+
+  private static final String RESOURCE_PATH = "/data/local/tmp/llama/";
+  private static final String TOKENIZER_PATH = "tokenizer.model";
+  private static final String MODEL_PATH = "stories110M.pte";
+
+  private final List<String> results = new ArrayList<>();
+
+  @Test
+  public void testLoadAndGenerate() {
+    String tokenizerPath = RESOURCE_PATH + TOKENIZER_PATH;
+    File model = new File(RESOURCE_PATH + MODEL_PATH);
+    LlmModule mModule = new LlmModule(model.getPath(), tokenizerPath, 0.8f);
+
+    int loadResult = mModule.load();
+    // Check that the model can be loaded successfully
+    assertEquals(0, loadResult);
+
+    // Run a testing prompt
+    mModule.generate("How do you do! I'm testing llm on mobile device", SanityCheck.this);
+
+    // Verify we got some response
+    assertFalse("Should receive at least one result token", results.isEmpty());
+  }
+
+  @Override
+  public void onResult(String result) {
+    results.add(result);
+  }
+
+  @Override
+  public void onStats(String result) {
+    // Not measuring performance for now
+  }
+}