Add a test for compos_key_service.

Also disabled testOdrefresh, until the memory size issue is resolved.

Add both tests to TEST_MAPPING.

Also gratuitously changed a few things.

Bug: 193603140
Test: atest ComposHostTestCases
Change-Id: I655153fda6792df5f2d0210dedb12fc3ab692af7
diff --git a/TEST_MAPPING b/TEST_MAPPING
index 69d4568..d8b294b 100644
--- a/TEST_MAPPING
+++ b/TEST_MAPPING
@@ -2,7 +2,11 @@
   "presubmit": [
     {
       "name": "MicrodroidHostTestCases"
+    },
+    {
+      "name": "ComposHostTestCases"
     }
+
   ],
   "postsubmit": [
     // TODO(jiyong): promote this to presubmit. That currently doesn't work because
diff --git a/compos/tests/java/android/compos/test/ComposKeyTestCase.java b/compos/tests/java/android/compos/test/ComposKeyTestCase.java
new file mode 100644
index 0000000..654dc0b
--- /dev/null
+++ b/compos/tests/java/android/compos/test/ComposKeyTestCase.java
@@ -0,0 +1,150 @@
+/*
+ * Copyright (C) 2021 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package android.compos.test;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import android.platform.test.annotations.RootPermissionTest;
+import android.virt.test.CommandRunner;
+import android.virt.test.VirtualizationTestCaseBase;
+
+import com.android.compatibility.common.util.PollingCheck;
+import com.android.tradefed.testtype.DeviceJUnit4ClassRunner;
+import com.android.tradefed.util.CommandResult;
+import com.android.tradefed.util.CommandStatus;
+
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+@RootPermissionTest
+@RunWith(DeviceJUnit4ClassRunner.class)
+public final class ComposKeyTestCase extends VirtualizationTestCaseBase {
+
+    /** Wait time for service to be ready on boot */
+    private static final int READY_LATENCY_MS = 10 * 1000; // 10 seconds
+
+    // Path to compos_key_cmd tool
+    private static final String COMPOS_KEY_CMD_BIN = "/apex/com.android.compos/bin/compos_key_cmd";
+
+    private String mCid;
+
+    @Before
+    public void setUp() throws Exception {
+        testIfDeviceIsCapable(getDevice());
+
+        prepareVirtualizationTestSetup(getDevice());
+    }
+
+    @After
+    public void tearDown() throws Exception {
+        if (mCid != null) {
+            shutdownMicrodroid(getDevice(), mCid);
+            mCid = null;
+        }
+
+        cleanUpVirtualizationTestSetup(getDevice());
+    }
+
+    @Test
+    public void testKeyService() throws Exception {
+        startVm();
+        waitForServiceRunning();
+
+        CommandRunner android = new CommandRunner(getDevice());
+        CommandResult result;
+
+        // Generate keys - should succeed
+        android.run(
+                COMPOS_KEY_CMD_BIN,
+                "--cid " + mCid,
+                "generate",
+                TEST_ROOT + "test_key.blob",
+                TEST_ROOT + "test_key.pubkey");
+
+        // Verify them - should also succeed, since we just generated them
+        android.run(
+                COMPOS_KEY_CMD_BIN,
+                "--cid " + mCid,
+                "verify",
+                TEST_ROOT + "test_key.blob",
+                TEST_ROOT + "test_key.pubkey");
+
+        // Swap public key & blob - should fail to verify
+        result =
+                android.runForResult(
+                        COMPOS_KEY_CMD_BIN,
+                        "--cid " + mCid,
+                        "verify",
+                        TEST_ROOT + "test_key.pubkey",
+                        TEST_ROOT + "test_key.blob");
+        assertThat(result.getStatus()).isEqualTo(CommandStatus.FAILED);
+
+        // Generate another set of keys - should succeed
+        android.run(
+                COMPOS_KEY_CMD_BIN,
+                "--cid " + mCid,
+                "generate",
+                TEST_ROOT + "test_key2.blob",
+                TEST_ROOT + "test_key2.pubkey");
+
+        // They should also verify ok
+        android.run(
+                COMPOS_KEY_CMD_BIN,
+                "--cid " + mCid,
+                "verify",
+                TEST_ROOT + "test_key2.blob",
+                TEST_ROOT + "test_key2.pubkey");
+
+        // Mismatched key blob & public key should fail to verify
+        result =
+                android.runForResult(
+                        COMPOS_KEY_CMD_BIN,
+                        "--cid " + mCid,
+                        "verify",
+                        TEST_ROOT + "test_key.pubkey",
+                        TEST_ROOT + "test_key2.blob");
+        assertThat(result.getStatus()).isEqualTo(CommandStatus.FAILED);
+    }
+
+    private void startVm() throws Exception {
+        final String apkName = "CompOSPayloadApp.apk";
+        final String packageName = "com.android.compos.payload";
+        mCid =
+                startMicrodroid(
+                        getDevice(),
+                        getBuild(),
+                        apkName,
+                        packageName,
+                        "assets/key_service_vm_config.json",
+                        /* debug */ true);
+        adbConnectToMicrodroid(getDevice(), mCid);
+    }
+
+    private void waitForServiceRunning() {
+        try {
+            PollingCheck.waitFor(READY_LATENCY_MS, this::isServiceRunning);
+        } catch (Exception e) {
+            throw new RuntimeException("Service unavailable", e);
+        }
+    }
+
+    private boolean isServiceRunning() {
+        return tryRunOnMicrodroid("pidof compos_key_main") != null;
+    }
+}
diff --git a/compos/tests/java/android/compos/test/ComposTestCase.java b/compos/tests/java/android/compos/test/ComposTestCase.java
index 503e0f9..72eeb2b 100644
--- a/compos/tests/java/android/compos/test/ComposTestCase.java
+++ b/compos/tests/java/android/compos/test/ComposTestCase.java
@@ -16,20 +16,19 @@
 
 package android.compos.test;
 
-import static org.hamcrest.CoreMatchers.is;
-import static org.junit.Assert.assertThat;
+import static com.google.common.truth.Truth.assertThat;
 
 import android.platform.test.annotations.RootPermissionTest;
 import android.virt.test.CommandRunner;
 import android.virt.test.VirtualizationTestCaseBase;
 
 import com.android.compatibility.common.util.PollingCheck;
-import com.android.tradefed.device.DeviceNotAvailableException;
 import com.android.tradefed.testtype.DeviceJUnit4ClassRunner;
 import com.android.tradefed.util.CommandResult;
 
 import org.junit.After;
 import org.junit.Before;
+import org.junit.Ignore;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
@@ -53,27 +52,16 @@
     private String mCid;
 
     @Before
-    public void setUp() throws DeviceNotAvailableException {
+    public void setUp() throws Exception {
         testIfDeviceIsCapable(getDevice());
 
         prepareVirtualizationTestSetup(getDevice());
 
-        final String apkName = "CompOSPayloadApp.apk";
-        final String packageName = "com.android.compos.payload";
-        final String configPath = "assets/vm_config.json"; // path inside the APK
-        mCid =
-                startMicrodroid(
-                        getDevice(),
-                        getBuild(),
-                        apkName,
-                        packageName,
-                        configPath,
-                        /* debug */ true);
-        adbConnectToMicrodroid(getDevice(), mCid);
+        startComposVm();
     }
 
     @After
-    public void tearDown() throws DeviceNotAvailableException {
+    public void tearDown() throws Exception {
         if (mCid != null) {
             shutdownMicrodroid(getDevice(), mCid);
             mCid = null;
@@ -83,7 +71,8 @@
     }
 
     @Test
-    public void testOdrefresh() throws DeviceNotAvailableException, InterruptedException {
+    @Ignore("b/192294431")
+    public void testOdrefresh() throws Exception {
         waitForServiceRunning();
 
         CommandRunner android = new CommandRunner(getDevice());
@@ -95,16 +84,30 @@
                         ODREFRESH_BIN,
                         "--use-compilation-os=" + mCid,
                         "--force-compile");
-        assertThat(result.getExitCode(), is(COMPILATION_SUCCESS));
+        assertThat(result.getExitCode()).isEqualTo(COMPILATION_SUCCESS);
 
         // Expect the output to be valid.
         result = android.runForResultWithTimeout(ODREFRESH_TIMEOUT_MS, ODREFRESH_BIN, "--check");
-        assertThat(result.getExitCode(), is(OKAY));
+        assertThat(result.getExitCode()).isEqualTo(OKAY);
+    }
+
+    private void startComposVm() throws Exception {
+        final String apkName = "CompOSPayloadApp.apk";
+        final String packageName = "com.android.compos.payload";
+        mCid =
+                startMicrodroid(
+                        getDevice(),
+                        getBuild(),
+                        apkName,
+                        packageName,
+                        "assets/vm_config.json",
+                        /* debug */ true);
+        adbConnectToMicrodroid(getDevice(), mCid);
     }
 
     private void waitForServiceRunning() {
         try {
-            PollingCheck.waitFor(COMPSVC_READY_LATENCY_MS, () -> isServiceRunning());
+            PollingCheck.waitFor(COMPSVC_READY_LATENCY_MS, this::isServiceRunning);
         } catch (Exception e) {
             throw new RuntimeException("Service unavailable", e);
         }
diff --git a/tests/hostside/helper/java/android/virt/test/VirtualizationTestCaseBase.java b/tests/hostside/helper/java/android/virt/test/VirtualizationTestCaseBase.java
index fef8864..4cb9dc0 100644
--- a/tests/hostside/helper/java/android/virt/test/VirtualizationTestCaseBase.java
+++ b/tests/hostside/helper/java/android/virt/test/VirtualizationTestCaseBase.java
@@ -40,7 +40,7 @@
 import java.util.regex.Pattern;
 
 public abstract class VirtualizationTestCaseBase extends BaseHostJUnit4Test {
-    private static final String TEST_ROOT = "/data/local/tmp/virt/";
+    protected static final String TEST_ROOT = "/data/local/tmp/virt/";
     private static final String VIRT_APEX = "/apex/com.android.virt/";
     private static final int TEST_VM_ADB_PORT = 8000;
     private static final String MICRODROID_SERIAL = "localhost:" + TEST_VM_ADB_PORT;