[TF] Support Sharding feature in ExecutableTargetTest

Bug: 161196833
Test: 1. unnittest
      2. atest vts_ltp_test_arm_64 --sharding
Change-Id: I73fc3cbefd601cd2091fb1c8672b39d98584a2ff
diff --git a/test_framework/com/android/tradefed/testtype/binary/ExecutableBaseTest.java b/test_framework/com/android/tradefed/testtype/binary/ExecutableBaseTest.java
index ed627bc..94e4aa6 100644
--- a/test_framework/com/android/tradefed/testtype/binary/ExecutableBaseTest.java
+++ b/test_framework/com/android/tradefed/testtype/binary/ExecutableBaseTest.java
@@ -15,6 +15,7 @@
  */
 package com.android.tradefed.testtype.binary;
 
+import com.google.common.annotations.VisibleForTesting;
 import com.android.tradefed.config.Option;
 import com.android.tradefed.config.OptionCopier;
 import com.android.tradefed.device.DeviceNotAvailableException;
@@ -89,6 +90,16 @@
     private Set<String> mIncludeFilters = new LinkedHashSet<>();
     private Set<String> mExcludeFilters = new LinkedHashSet<>();
 
+    /**
+     * Get test commands.
+     *
+     * @return the test commands.
+     */
+    @VisibleForTesting
+    Map<String, String> getTestCommands() {
+        return mTestCommands;
+    }
+
     /** @return the timeout applied to each binary for their execution. */
     protected long getTimeoutPerBinaryMs() {
         return mTimeoutPerBinaryMs;
@@ -254,26 +265,44 @@
     /** {@inheritDoc} */
     @Override
     public final Collection<IRemoteTest> split() {
-        if (mBinaryPaths.size() <= 2) {
+        int testCount = mBinaryPaths.size() + mTestCommands.size();
+        if (testCount <= 2) {
             return null;
         }
         Collection<IRemoteTest> tests = new ArrayList<>();
         for (String path : mBinaryPaths) {
-            tests.add(getTestShard(path));
+            tests.add(getTestShard(path, null, null));
+        }
+        Map<String, String> testCommands = new LinkedHashMap<>(mTestCommands);
+        for (String testName : testCommands.keySet()) {
+            String cmd = testCommands.get(testName);
+            tests.add(getTestShard(null, testName, cmd));
         }
         return tests;
     }
 
-    private IRemoteTest getTestShard(String path) {
+    /**
+     * Get a testShard of ExecutableBaseTest.
+     *
+     * @param binaryPath the binary path for ExecutableHostTest.
+     * @param testName the test name for ExecutableTargetTest.
+     * @param cmd the test command for ExecutableTargetTest.
+     * @return a shard{@link IRemoteTest} of ExecutableBaseTest{@link ExecutableBaseTest}
+     */
+    private IRemoteTest getTestShard(String binaryPath, String testName, String cmd) {
         ExecutableBaseTest shard = null;
         try {
             shard = this.getClass().getDeclaredConstructor().newInstance();
             OptionCopier.copyOptionsNoThrow(this, shard);
-            // We approximate the runtime of each shard to be equal since we can't know.
-            shard.mRuntimeHintMs = mRuntimeHintMs / shard.mBinaryPaths.size();
-            // Set one binary per shard
             shard.mBinaryPaths.clear();
-            shard.mBinaryPaths.add(path);
+            shard.mTestCommands.clear();
+            if (binaryPath != null) {
+                // Set one binary per shard
+                shard.mBinaryPaths.add(binaryPath);
+            } else if (testName != null && cmd != null) {
+                // Set one test command per shard
+                shard.mTestCommands.put(testName, cmd);
+            }
         } catch (InstantiationException
                 | IllegalAccessException
                 | InvocationTargetException
diff --git a/tests/src/com/android/tradefed/testtype/binary/ExecutableTargetTestTest.java b/tests/src/com/android/tradefed/testtype/binary/ExecutableTargetTestTest.java
index 4c0f9bd..ba18dc0 100644
--- a/tests/src/com/android/tradefed/testtype/binary/ExecutableTargetTestTest.java
+++ b/tests/src/com/android/tradefed/testtype/binary/ExecutableTargetTestTest.java
@@ -15,6 +15,7 @@
  */
 package com.android.tradefed.testtype.binary;
 
+import static org.junit.Assert.assertEquals;
 import static org.mockito.ArgumentMatchers.eq;
 
 import com.android.tradefed.config.ConfigurationException;
@@ -29,6 +30,7 @@
 import com.android.tradefed.result.TestDescription;
 import com.android.tradefed.result.error.InfraErrorIdentifier;
 import com.android.tradefed.result.proto.TestRecordProto.FailureStatus;
+import com.android.tradefed.testtype.IRemoteTest;
 import com.android.tradefed.util.CommandResult;
 
 import org.junit.Before;
@@ -38,6 +40,8 @@
 import org.mockito.Mockito;
 
 import java.util.HashMap;
+import java.util.Collection;
+import java.util.Map;
 
 /** Unit tests for {@link com.android.tradefed.testtype.binary.ExecutableTargetTest}. */
 @RunWith(JUnit4.class)
@@ -369,4 +373,31 @@
                         Mockito.eq(testDescription3),
                         Mockito.eq(new HashMap<String, MetricMeasurement.Metric>()));
     }
+
+    /** Test split() for sharding */
+    @Test
+    public void testShard_Split() throws DeviceNotAvailableException, ConfigurationException {
+        mExecutableTargetTest = new ExecutableTargetTest();
+        // Set test commands
+        OptionSetter setter = new OptionSetter(mExecutableTargetTest);
+        setter.setOptionValue("test-command-line", testName1, testCmd1);
+        setter.setOptionValue("test-command-line", testName2, testCmd2);
+        setter.setOptionValue("test-command-line", testName3, testCmd3);
+        // Split the shard.
+        Collection<IRemoteTest> testShards = mExecutableTargetTest.split();
+        // Test the size of the test Shard.
+        assertEquals(3, testShards.size());
+        // Test the command of each shard.
+        for (IRemoteTest test : testShards) {
+            Map<String, String> TestCommands = ((ExecutableTargetTest) test).getTestCommands();
+            String cmd1 = TestCommands.get(testName1);
+            if (cmd1 != null) assertEquals(testCmd1, cmd1);
+            String cmd2 = TestCommands.get(testName2);
+            if (cmd2 != null) assertEquals(testCmd2, cmd2);
+            String cmd3 = TestCommands.get(testName3);
+            if (cmd3 != null) assertEquals(testCmd3, cmd3);
+            // The test command should equals to one of them.
+            assertEquals(true, cmd1 != null || cmd2 != null || cmd3 != null);
+        }
+    }
 }