release-request-6cde801d-7c1b-44a9-92d6-15dfa0829dd6-for-git_oc-vts-release-4243288 snap-temp-L53400000089083260

Change-Id: Ie5b8b72ce5b336edc3de0a0b89b1c8df34fc8e49
diff --git a/src/com/android/tradefed/device/LargeOutputReceiver.java b/src/com/android/tradefed/device/LargeOutputReceiver.java
index 8af95af..d402db1 100644
--- a/src/com/android/tradefed/device/LargeOutputReceiver.java
+++ b/src/com/android/tradefed/device/LargeOutputReceiver.java
@@ -82,7 +82,7 @@
     public synchronized InputStreamSource getData() {
         if (mOutStream != null) {
             try {
-                return new SnapshotInputStreamSource(mOutStream.getData());
+                return new SnapshotInputStreamSource("LargeOutputReceiver", mOutStream.getData());
             } catch (IOException e) {
                 CLog.e("failed to get %s data for %s.", mDescriptor, mSerialNumber);
                 CLog.e(e);
diff --git a/src/com/android/tradefed/device/NativeDevice.java b/src/com/android/tradefed/device/NativeDevice.java
index cf3bbbb..ca62f2b 100644
--- a/src/com/android/tradefed/device/NativeDevice.java
+++ b/src/com/android/tradefed/device/NativeDevice.java
@@ -3177,7 +3177,8 @@
                         getSerialNumber());
             } else {
                 try {
-                    return new SnapshotInputStreamSource(mEmulatorOutput.getData());
+                    return new SnapshotInputStreamSource(
+                            "getEmulatorOutput", mEmulatorOutput.getData());
                 } catch (IOException e) {
                     CLog.e("Failed to get %s data.", getSerialNumber());
                     CLog.e(e);
diff --git a/src/com/android/tradefed/invoker/shard/StrictShardHelper.java b/src/com/android/tradefed/invoker/shard/StrictShardHelper.java
index 7ea20b8..c2e5043 100644
--- a/src/com/android/tradefed/invoker/shard/StrictShardHelper.java
+++ b/src/com/android/tradefed/invoker/shard/StrictShardHelper.java
@@ -25,6 +25,7 @@
 import com.android.tradefed.testtype.IShardableTest;
 import com.android.tradefed.testtype.IStrictShardableTest;
 import com.android.tradefed.testtype.suite.ITestSuite;
+import com.android.tradefed.testtype.suite.ModuleMerger;
 
 import java.util.ArrayList;
 import java.util.Collection;
@@ -54,10 +55,10 @@
         } else {
             List<IRemoteTest> listAllTests = getAllTests(config, shardCount, context);
             // We cannot shuffle to get better average results
-            // TODO: normalize the distribution of tests: current distributions is pretty much
-            // the order it was split into which tend to be unbalanced
             normalizeDistribution(listAllTests, shardCount);
-            config.setTests(splitTests(listAllTests, shardCount, shardIndex));
+            List<IRemoteTest> splitList = splitTests(listAllTests, shardCount, shardIndex);
+            aggregateSuiteModules(splitList);
+            config.setTests(splitList);
         }
         return false;
     }
@@ -170,4 +171,32 @@
             }
         }
     }
+
+    /**
+     * Special handling for suite from {@link ITestSuite}. We aggregate the tests in the same shard
+     * in order to optimize target_preparation step.
+     *
+     * @param tests the {@link List} of {@link IRemoteTest} for that shard.
+     */
+    private void aggregateSuiteModules(List<IRemoteTest> tests) {
+        List<IRemoteTest> dupList = new ArrayList<>(tests);
+        for (int i = 0; i < dupList.size(); i++) {
+            if (dupList.get(i) instanceof ITestSuite) {
+                // We iterate the other tests to see if we can find another from the same module.
+                for (int j = i + 1; j < dupList.size(); j++) {
+                    // If the test was not already merged
+                    if (tests.contains(dupList.get(j))) {
+                        if (dupList.get(j) instanceof ITestSuite) {
+                            if (ModuleMerger.arePartOfSameSuite(
+                                    (ITestSuite) dupList.get(i), (ITestSuite) dupList.get(j))) {
+                                ModuleMerger.mergeSplittedITestSuite(
+                                        (ITestSuite) dupList.get(i), (ITestSuite) dupList.get(j));
+                                tests.remove(dupList.get(j));
+                            }
+                        }
+                    }
+                }
+            }
+        }
+    }
 }
diff --git a/src/com/android/tradefed/log/FileLogger.java b/src/com/android/tradefed/log/FileLogger.java
index 12d3e24..f9b13d9 100644
--- a/src/com/android/tradefed/log/FileLogger.java
+++ b/src/com/android/tradefed/log/FileLogger.java
@@ -197,7 +197,7 @@
             try {
                 // create a InputStream from log file
                 mLogStream.flush();
-                return new SnapshotInputStreamSource(mLogStream.getData());
+                return new SnapshotInputStreamSource("FileLogger", mLogStream.getData());
             } catch (IOException e) {
                 System.err.println("Failed to get log");
                 e.printStackTrace();
diff --git a/src/com/android/tradefed/result/SnapshotInputStreamSource.java b/src/com/android/tradefed/result/SnapshotInputStreamSource.java
index 464d927..cc9f803 100644
--- a/src/com/android/tradefed/result/SnapshotInputStreamSource.java
+++ b/src/com/android/tradefed/result/SnapshotInputStreamSource.java
@@ -32,16 +32,14 @@
     private File mBackingFile;
     private boolean mIsCancelled = false;
 
-    /**
-     * Constructor for a file-backed {@link InputStreamSource}
-     */
-    public SnapshotInputStreamSource(InputStream stream) {
+    /** Constructor for a file-backed {@link InputStreamSource} */
+    public SnapshotInputStreamSource(String name, InputStream stream) {
         if (stream == null) {
             throw new NullPointerException();
         }
 
         try {
-            mBackingFile = createBackingFile(stream);
+            mBackingFile = createBackingFile(name, stream);
         } catch (IOException e) {
             // Log an error and invalidate ourself
             CLog.e("Received IOException while trying to wrap a stream");
@@ -52,11 +50,12 @@
 
     /**
      * Create the backing file and fill it with the contents of {@code stream}.
-     * <p />
-     * Exposed for unit testing
+     *
+     * <p>Exposed for unit testing
      */
-    File createBackingFile(InputStream stream) throws IOException {
-        File backingFile = FileUtil.createTempFile(this.getClass().getSimpleName() + "_", ".txt");
+    File createBackingFile(String name, InputStream stream) throws IOException {
+        File backingFile =
+                FileUtil.createTempFile(name + "_" + this.getClass().getSimpleName() + "_", ".txt");
         FileUtil.writeToFile(stream, backingFile);
         return backingFile;
     }
diff --git a/src/com/android/tradefed/testtype/suite/ITestSuite.java b/src/com/android/tradefed/testtype/suite/ITestSuite.java
index f5a1019..d6c89cb 100644
--- a/src/com/android/tradefed/testtype/suite/ITestSuite.java
+++ b/src/com/android/tradefed/testtype/suite/ITestSuite.java
@@ -478,4 +478,12 @@
     public void setInvocationContext(IInvocationContext invocationContext) {
         mContext = invocationContext;
     }
+
+    /**
+     * Returns the {@link ModuleDefinition} to be executed directly, or null if none yet (when the
+     * ITestSuite has not been sharded yet).
+     */
+    public ModuleDefinition getDirectModule() {
+        return mDirectModule;
+    }
 }
diff --git a/src/com/android/tradefed/testtype/suite/ModuleDefinition.java b/src/com/android/tradefed/testtype/suite/ModuleDefinition.java
index daf4bb1..ce89cb2 100644
--- a/src/com/android/tradefed/testtype/suite/ModuleDefinition.java
+++ b/src/com/android/tradefed/testtype/suite/ModuleDefinition.java
@@ -139,6 +139,23 @@
     }
 
     /**
+     * Add some {@link IRemoteTest} to be executed as part of the module. Used when merging two
+     * modules.
+     */
+    void addTests(List<IRemoteTest> test) {
+        synchronized (mTests) {
+            mTests.addAll(test);
+        }
+    }
+
+    /** Returns the current number of {@link IRemoteTest} waiting to be executed. */
+    public int numTests() {
+        synchronized (mTests) {
+            return mTests.size();
+        }
+    }
+
+    /**
      * Return True if the Module still has {@link IRemoteTest} to run in its pool. False otherwise.
      */
     protected boolean hasTests() {
diff --git a/src/com/android/tradefed/testtype/suite/ModuleMerger.java b/src/com/android/tradefed/testtype/suite/ModuleMerger.java
new file mode 100644
index 0000000..9410a2e
--- /dev/null
+++ b/src/com/android/tradefed/testtype/suite/ModuleMerger.java
@@ -0,0 +1,64 @@
+/*
+ * Copyright (C) 2017 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 com.android.tradefed.testtype.suite;
+
+/**
+ * Helper class for operation related to merging {@link ITestSuite} and {@link ModuleDefinition}
+ * after a split.
+ */
+public class ModuleMerger {
+
+    private static void mergeModules(ModuleDefinition module1, ModuleDefinition module2) {
+        if (!module1.getId().equals(module2.getId())) {
+            throw new IllegalArgumentException(
+                    String.format(
+                            "Modules must have the same id to be mergeable: received %s and "
+                                    + "%s",
+                            module1.getId(), module2.getId()));
+        }
+        module1.addTests(module2.getTests());
+    }
+
+    /**
+     * Merge the modules from one suite to another.
+     *
+     * @param suite1 the suite that will receive the module from the other.
+     * @param suite2 the suite that will give the module.
+     */
+    public static void mergeSplittedITestSuite(ITestSuite suite1, ITestSuite suite2) {
+        if (suite1.getDirectModule() == null) {
+            throw new IllegalArgumentException("suite was not a splitted suite.");
+        }
+        if (suite2.getDirectModule() == null) {
+            throw new IllegalArgumentException("suite was not a splitted suite.");
+        }
+        mergeModules(suite1.getDirectModule(), suite2.getDirectModule());
+    }
+
+    /** Returns true if the two suites are part of the same original split. False otherwise. */
+    public static boolean arePartOfSameSuite(ITestSuite suite1, ITestSuite suite2) {
+        if (suite1.getDirectModule() == null) {
+            return false;
+        }
+        if (suite2.getDirectModule() == null) {
+            return false;
+        }
+        if (!suite1.getDirectModule().getId().equals(suite2.getDirectModule().getId())) {
+            return false;
+        }
+        return true;
+    }
+}
diff --git a/tests/src/com/android/tradefed/UnitTests.java b/tests/src/com/android/tradefed/UnitTests.java
index 62ff479..d654f01 100644
--- a/tests/src/com/android/tradefed/UnitTests.java
+++ b/tests/src/com/android/tradefed/UnitTests.java
@@ -152,6 +152,7 @@
 import com.android.tradefed.testtype.suite.ITestSuiteTest;
 import com.android.tradefed.testtype.suite.ModuleDefinitionTest;
 import com.android.tradefed.testtype.suite.ModuleListenerTest;
+import com.android.tradefed.testtype.suite.ModuleMergerTest;
 import com.android.tradefed.testtype.suite.ModuleSplitterTest;
 import com.android.tradefed.testtype.suite.TestFailureListenerTest;
 import com.android.tradefed.testtype.suite.TfSuiteRunnerTest;
@@ -394,6 +395,7 @@
     ITestSuiteTest.class,
     ModuleDefinitionTest.class,
     ModuleListenerTest.class,
+    ModuleMergerTest.class,
     ModuleSplitterTest.class,
     TestFailureListenerTest.class,
     TfSuiteRunnerTest.class,
diff --git a/tests/src/com/android/tradefed/invoker/shard/StrictShardHelperTest.java b/tests/src/com/android/tradefed/invoker/shard/StrictShardHelperTest.java
index 29f0ac4..1d46c84 100644
--- a/tests/src/com/android/tradefed/invoker/shard/StrictShardHelperTest.java
+++ b/tests/src/com/android/tradefed/invoker/shard/StrictShardHelperTest.java
@@ -15,14 +15,21 @@
  */
 package com.android.tradefed.invoker.shard;
 
-import static org.junit.Assert.*;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotEquals;
+import static org.junit.Assert.assertSame;
+import static org.junit.Assert.assertTrue;
 
 import com.android.tradefed.build.BuildInfo;
 import com.android.tradefed.command.CommandOptions;
 import com.android.tradefed.config.Configuration;
+import com.android.tradefed.config.ConfigurationException;
+import com.android.tradefed.config.ConfigurationFactory;
 import com.android.tradefed.config.IConfiguration;
 import com.android.tradefed.config.OptionSetter;
 import com.android.tradefed.device.DeviceNotAvailableException;
+import com.android.tradefed.device.ITestDevice;
 import com.android.tradefed.invoker.IInvocationContext;
 import com.android.tradefed.invoker.IRescheduler;
 import com.android.tradefed.invoker.InvocationContext;
@@ -30,7 +37,9 @@
 import com.android.tradefed.result.ITestInvocationListener;
 import com.android.tradefed.testtype.IRemoteTest;
 import com.android.tradefed.testtype.StubTest;
+import com.android.tradefed.testtype.suite.ITestSuite;
 
+import org.easymock.EasyMock;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -38,6 +47,10 @@
 import org.mockito.ArgumentMatcher;
 import org.mockito.Mockito;
 
+import java.util.ArrayList;
+import java.util.LinkedHashMap;
+import java.util.List;
+
 /** Unit tests for {@link StrictShardHelper}. */
 @RunWith(JUnit4.class)
 public class StrictShardHelperTest {
@@ -173,4 +186,109 @@
         // We have no tests to put in shard-index 1 so it's empty.
         assertEquals(0, mConfig.getTests().size());
     }
+
+    /** Test class to simulate an ITestSuite getting split. */
+    public static class SplitITestSuite extends ITestSuite {
+
+        private String mName;
+
+        public SplitITestSuite() {}
+
+        public SplitITestSuite(String name) {
+            mName = name;
+        }
+
+        @Override
+        public LinkedHashMap<String, IConfiguration> loadTests() {
+            LinkedHashMap<String, IConfiguration> configs = new LinkedHashMap<>();
+            IConfiguration configuration = null;
+            try {
+                configuration =
+                        ConfigurationFactory.getInstance()
+                                .createConfigurationFromArgs(
+                                        new String[] {"empty", "--num-shards", "2"});
+            } catch (ConfigurationException e) {
+                throw new RuntimeException(e);
+            }
+            configs.put(mName, configuration);
+            return configs;
+        }
+    }
+
+    private ITestSuite createFakeSuite(String name) throws Exception {
+        ITestSuite suite = new SplitITestSuite(name);
+        return suite;
+    }
+
+    private List<IRemoteTest> testShard(int shardIndex) throws Exception {
+        mContext.addAllocatedDevice("default", EasyMock.createMock(ITestDevice.class));
+        List<IRemoteTest> test = new ArrayList<>();
+        test.add(createFakeSuite("module2"));
+        test.add(createFakeSuite("module1"));
+        test.add(createFakeSuite("module3"));
+        test.add(createFakeSuite("module1"));
+        test.add(createFakeSuite("module1"));
+        test.add(createFakeSuite("module2"));
+        test.add(createFakeSuite("module3"));
+        CommandOptions options = new CommandOptions();
+        OptionSetter setter = new OptionSetter(options);
+        setter.setOptionValue("disable-strict-sharding", "true");
+        setter.setOptionValue("shard-count", "3");
+        setter.setOptionValue("shard-index", Integer.toString(shardIndex));
+        mConfig.setCommandOptions(options);
+        mConfig.setCommandLine(new String[] {"empty"});
+        mConfig.setTests(test);
+        mHelper.shardConfig(mConfig, mContext, mRescheduler);
+        return mConfig.getTests();
+    }
+
+    /**
+     * Total for all the _shardX test should be 14 tests (2 per modules). 6 for module1: 3 module1
+     * shard * 2 4 for module2: 2 module2 shard * 2 4 for module3: 2 module3 shard * 2
+     */
+    @Test
+    public void testMergeSuite_shard0() throws Exception {
+        List<IRemoteTest> res = testShard(0);
+        assertEquals(3, res.size());
+
+        assertTrue(res.get(0) instanceof ITestSuite);
+        assertEquals("module1", ((ITestSuite) res.get(0)).getDirectModule().getId());
+        assertEquals(3, ((ITestSuite) res.get(0)).getDirectModule().numTests());
+
+        assertTrue(res.get(1) instanceof ITestSuite);
+        assertEquals("module3", ((ITestSuite) res.get(1)).getDirectModule().getId());
+        assertEquals(1, ((ITestSuite) res.get(1)).getDirectModule().numTests());
+
+        assertTrue(res.get(2) instanceof ITestSuite);
+        assertEquals("module2", ((ITestSuite) res.get(2)).getDirectModule().getId());
+        assertEquals(1, ((ITestSuite) res.get(2)).getDirectModule().numTests());
+    }
+
+    @Test
+    public void testMergeSuite_shard1() throws Exception {
+        List<IRemoteTest> res = testShard(1);
+        assertEquals(2, res.size());
+
+        assertTrue(res.get(0) instanceof ITestSuite);
+        assertEquals("module3", ((ITestSuite) res.get(0)).getDirectModule().getId());
+        assertEquals(2, ((ITestSuite) res.get(0)).getDirectModule().numTests());
+
+        assertTrue(res.get(1) instanceof ITestSuite);
+        assertEquals("module2", ((ITestSuite) res.get(1)).getDirectModule().getId());
+        assertEquals(3, ((ITestSuite) res.get(1)).getDirectModule().numTests());
+    }
+
+    @Test
+    public void testMergeSuite_shard2() throws Exception {
+        List<IRemoteTest> res = testShard(2);
+        assertEquals(2, res.size());
+
+        assertTrue(res.get(0) instanceof ITestSuite);
+        assertEquals("module1", ((ITestSuite) res.get(0)).getDirectModule().getId());
+        assertEquals(3, ((ITestSuite) res.get(0)).getDirectModule().numTests());
+
+        assertTrue(res.get(1) instanceof ITestSuite);
+        assertEquals("module3", ((ITestSuite) res.get(1)).getDirectModule().getId());
+        assertEquals(1, ((ITestSuite) res.get(1)).getDirectModule().numTests());
+    }
 }
diff --git a/tests/src/com/android/tradefed/result/SnapshotInputStreamSourceTest.java b/tests/src/com/android/tradefed/result/SnapshotInputStreamSourceTest.java
index 0c5a559..943d400 100644
--- a/tests/src/com/android/tradefed/result/SnapshotInputStreamSourceTest.java
+++ b/tests/src/com/android/tradefed/result/SnapshotInputStreamSourceTest.java
@@ -47,12 +47,13 @@
             }
         };
 
-        InputStreamSource source = new SnapshotInputStreamSource(mInputStream) {
-            @Override
-            File createBackingFile(InputStream stream) {
-                return fakeFile;
-            }
-        };
+        InputStreamSource source =
+                new SnapshotInputStreamSource("SnapUnitTest", mInputStream) {
+                    @Override
+                    File createBackingFile(String name, InputStream stream) {
+                        return fakeFile;
+                    }
+                };
 
         try {
             source.cancel();
diff --git a/tests/src/com/android/tradefed/testtype/suite/ModuleMergerTest.java b/tests/src/com/android/tradefed/testtype/suite/ModuleMergerTest.java
new file mode 100644
index 0000000..931646c
--- /dev/null
+++ b/tests/src/com/android/tradefed/testtype/suite/ModuleMergerTest.java
@@ -0,0 +1,153 @@
+/*
+ * Copyright (C) 2017 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 com.android.tradefed.testtype.suite;
+
+import static org.junit.Assert.*;
+
+import com.android.tradefed.invoker.shard.StrictShardHelperTest.SplitITestSuite;
+import com.android.tradefed.testtype.IRemoteTest;
+
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+import java.util.Collection;
+import java.util.Iterator;
+
+/** Unit tests for {@link ModuleMerger}. */
+@RunWith(JUnit4.class)
+public class ModuleMergerTest {
+
+    /**
+     * Test that {@link ModuleMerger#arePartOfSameSuite(ITestSuite, ITestSuite)} returns false when
+     * the first suite is not splitted yet.
+     */
+    @Test
+    public void testPartOfSameSuite_notSplittedYet() {
+        SplitITestSuite suite1 = new SplitITestSuite("module1");
+        SplitITestSuite suite2 = new SplitITestSuite("module2");
+        assertFalse(ModuleMerger.arePartOfSameSuite(suite1, suite2));
+    }
+
+    /**
+     * Test that {@link ModuleMerger#arePartOfSameSuite(ITestSuite, ITestSuite)} returns false when
+     * the second suite is not splitted yet.
+     */
+    @Test
+    public void testPartOfSameSuite_notSplittedYet2() {
+        SplitITestSuite suite1 = new SplitITestSuite("module1");
+        Collection<IRemoteTest> res1 = suite1.split(2);
+        SplitITestSuite suite2 = new SplitITestSuite("module2");
+        assertFalse(ModuleMerger.arePartOfSameSuite((ITestSuite) res1.iterator().next(), suite2));
+    }
+
+    /**
+     * Test that {@link ModuleMerger#arePartOfSameSuite(ITestSuite, ITestSuite)} returns true when
+     * the two suites are splitted and from the same module.
+     */
+    @Test
+    public void testPartOfSameSuite_sameSuite() {
+        SplitITestSuite suite1 = new SplitITestSuite("module1");
+        Collection<IRemoteTest> res1 = suite1.split(2);
+        Iterator<IRemoteTest> ite = res1.iterator();
+        assertTrue(
+                ModuleMerger.arePartOfSameSuite((ITestSuite) ite.next(), (ITestSuite) ite.next()));
+    }
+
+    /**
+     * Test that {@link ModuleMerger#arePartOfSameSuite(ITestSuite, ITestSuite)} returns false when
+     * the two suites are splitted but from different modules.
+     */
+    @Test
+    public void testPartOfSameSuite_notSameSuite() {
+        SplitITestSuite suite1 = new SplitITestSuite("module1");
+        Collection<IRemoteTest> res1 = suite1.split(2);
+        SplitITestSuite suite2 = new SplitITestSuite("module2");
+        Collection<IRemoteTest> res2 = suite2.split(2);
+        assertFalse(
+                ModuleMerger.arePartOfSameSuite(
+                        (ITestSuite) res1.iterator().next(), (ITestSuite) res2.iterator().next()));
+    }
+
+    /**
+     * Test that {@link ModuleMerger#mergeSplittedITestSuite(ITestSuite, ITestSuite)} throws an
+     * exception when the first suite is not splitted yet.
+     */
+    @Test
+    public void testMergeSplittedITestSuite_notSplittedYet() {
+        SplitITestSuite suite1 = new SplitITestSuite("module1");
+        SplitITestSuite suite2 = new SplitITestSuite("module2");
+        try {
+            ModuleMerger.mergeSplittedITestSuite(suite1, suite2);
+            fail("Should have thrown an exception.");
+        } catch (IllegalArgumentException expected) {
+            // expected
+        }
+    }
+
+    /**
+     * Test that {@link ModuleMerger#mergeSplittedITestSuite(ITestSuite, ITestSuite)} throws an
+     * exception when the second suite is not splitted yet.
+     */
+    @Test
+    public void testMergeSplittedITestSuite_notSplittedYet2() {
+        SplitITestSuite suite1 = new SplitITestSuite("module1");
+        Collection<IRemoteTest> res1 = suite1.split(2);
+        SplitITestSuite suite2 = new SplitITestSuite("module2");
+        try {
+            ModuleMerger.mergeSplittedITestSuite((ITestSuite) res1.iterator().next(), suite2);
+            fail("Should have thrown an exception.");
+        } catch (IllegalArgumentException expected) {
+            // expected
+        }
+    }
+
+    /**
+     * Test that {@link ModuleMerger#mergeSplittedITestSuite(ITestSuite, ITestSuite)} throws an
+     * exception when the two suites are splitted but coming from different modules.
+     */
+    @Test
+    public void testMergeSplittedITestSuite_splittedSuiteFromDifferentModules() {
+        SplitITestSuite suite1 = new SplitITestSuite("module1");
+        Collection<IRemoteTest> res1 = suite1.split(2);
+        SplitITestSuite suite2 = new SplitITestSuite("module2");
+        Collection<IRemoteTest> res2 = suite2.split(2);
+        try {
+            ModuleMerger.mergeSplittedITestSuite(
+                    (ITestSuite) res1.iterator().next(), (ITestSuite) res2.iterator().next());
+            fail("Should have thrown an exception.");
+        } catch (IllegalArgumentException expected) {
+            // expected
+        }
+    }
+
+    /**
+     * Test that {@link ModuleMerger#mergeSplittedITestSuite(ITestSuite, ITestSuite)} properly
+     * assigns tests from the second suite to the first since part of same module.
+     */
+    @Test
+    public void testMergeSplittedITestSuite() {
+        SplitITestSuite suite1 = new SplitITestSuite("module1");
+        Collection<IRemoteTest> res1 = suite1.split(2);
+        Iterator<IRemoteTest> ite = res1.iterator();
+        ITestSuite split1 = (ITestSuite) ite.next();
+        ITestSuite split2 = (ITestSuite) ite.next();
+        assertEquals(2, split1.getDirectModule().numTests());
+        assertEquals(2, split2.getDirectModule().numTests());
+        ModuleMerger.mergeSplittedITestSuite(split1, split2);
+        assertEquals(4, split1.getDirectModule().numTests());
+    }
+}