Add sharding support for module generator

Previously the module generator is implemented as a target preparer. However, a target preparer is executed after the sharding process has finished. The only way to make the generated modules available for sharding without making changes to TradeFed's core code is to disguise this module generator as an instance of IShardableTest and declare it separately in test plan config. This is hacky, and in the long term a TradeFed centered solution is desired. For more details, see go/sharding-hack-for-module-gen.

Test: atest :presubmit
Test: Verified with sharding option in go/forrest-ui https://android-build.googleplex.com/builds/forrest/run/L36400000795517256
Fix: 176106448
Change-Id: I9b571ed9740796a4913179084f8e4a69d73f8c71
Merged-In: I9b571ed9740796a4913179084f8e4a69d73f8c71
diff --git a/harness/src/main/java/com/android/csuite/config/ModuleGenerator.java b/harness/src/main/java/com/android/csuite/config/ModuleGenerator.java
new file mode 100644
index 0000000..b37e7ef
--- /dev/null
+++ b/harness/src/main/java/com/android/csuite/config/ModuleGenerator.java
@@ -0,0 +1,220 @@
+/*
+ * Copyright (C) 2020 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.csuite.config;
+
+import com.android.compatibility.common.tradefed.build.CompatibilityBuildHelper;
+import com.android.tradefed.build.IBuildInfo;
+import com.android.tradefed.config.Option;
+import com.android.tradefed.config.Option.Importance;
+import com.android.tradefed.device.DeviceNotAvailableException;
+import com.android.tradefed.invoker.TestInformation;
+import com.android.tradefed.log.LogUtil.CLog;
+import com.android.tradefed.result.ITestInvocationListener;
+import com.android.tradefed.targetprep.ITargetPreparer;
+import com.android.tradefed.testtype.IBuildReceiver;
+import com.android.tradefed.testtype.IRemoteTest;
+import com.android.tradefed.testtype.IShardableTest;
+
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.io.Resources;
+
+import java.io.IOException;
+import java.nio.charset.StandardCharsets;
+import java.nio.file.FileSystem;
+import java.nio.file.FileSystems;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.Collection;
+import java.util.HashSet;
+import java.util.Set;
+
+/**
+ * A tool for generating TradeFed suite modules during runtime.
+ *
+ * <p>This class generates module config files into TradeFed's test directory at runtime using a
+ * template. Since the content of the test directory relies on what is being generated in a test
+ * run, there can only be one instance executing at a given time.
+ *
+ * <p>The intention of this class is to generate test modules at the beginning of a test run and
+ * cleans up after all tests finish, which resembles a target preparer. However, a target preparer
+ * is executed after the sharding process has finished. The only way to make the generated modules
+ * available for sharding without making changes to TradeFed's core code is to disguise this module
+ * generator as an instance of IShardableTest and declare it separately in test plan config. This is
+ * hacky, and in the long term a TradeFed centered solution is desired. For more details, see
+ * go/sharding-hack-for-module-gen. Note that since the generate step is executed as a test instance
+ * and cleanup step is executed as a target preparer, there should be no saved states between
+ * generating and cleaning up module files.
+ */
+public final class ModuleGenerator
+        implements IRemoteTest, IShardableTest, IBuildReceiver, ITargetPreparer {
+
+    @VisibleForTesting static final String MODULE_FILE_EXTENSION = ".config";
+    @VisibleForTesting static final String OPTION_TEMPLATE = "template";
+    @VisibleForTesting static final String OPTION_PACKAGE = "package";
+    private static final String TEMPLATE_PACKAGE_PATTERN = "\\{package\\}";
+    private static final Collection<IRemoteTest> NOT_SPLITABLE = null;
+
+    @Option(
+            name = OPTION_TEMPLATE,
+            description = "Module config template resource path.",
+            importance = Importance.ALWAYS)
+    private String mTemplate;
+
+    @Option(name = OPTION_PACKAGE, description = "App package names.")
+    private final Set<String> mPackages = new HashSet<>();
+
+    private final TestDirectoryProvider mTestDirectoryProvider;
+    private final ResourceLoader mResourceLoader;
+    private final FileSystem mFileSystem;
+    private IBuildInfo mBuildInfo;
+
+    public ModuleGenerator(Set<String> mPackages) {
+        this(FileSystems.getDefault());
+    }
+
+    public ModuleGenerator() {
+        this(FileSystems.getDefault());
+    }
+
+    private ModuleGenerator(FileSystem fileSystem) {
+        this(
+                fileSystem,
+                new CompatibilityTestDirectoryProvider(fileSystem),
+                new ClassResourceLoader());
+    }
+
+    @VisibleForTesting
+    ModuleGenerator(
+            FileSystem fileSystem,
+            TestDirectoryProvider testDirectoryProvider,
+            ResourceLoader resourceLoader) {
+        mFileSystem = fileSystem;
+        mTestDirectoryProvider = testDirectoryProvider;
+        mResourceLoader = resourceLoader;
+    }
+
+    @Override
+    public void run(final TestInformation testInfo, final ITestInvocationListener listener) {
+        // Intentionally left blank since this class is not really a test.
+    }
+
+    @Override
+    public void setUp(TestInformation testInfo) {
+        // Intentionally left blank.
+    }
+
+    @Override
+    public void setBuild(IBuildInfo buildInfo) {
+        mBuildInfo = buildInfo;
+    }
+
+    /**
+     * Generates test modules. Note that the implementation of this method is not related to
+     * sharding in any way.
+     */
+    @Override
+    public Collection<IRemoteTest> split() {
+        try {
+            // Executes the generate step.
+            generateModules();
+        } catch (IOException e) {
+            throw new RuntimeException("Failed to generate modules", e);
+        }
+
+        return NOT_SPLITABLE;
+    }
+
+    /** Cleans up generated test modules. */
+    @Override
+    public void tearDown(TestInformation testInfo, Throwable e) throws DeviceNotAvailableException {
+        // Gets build info from test info as when the class is executed as a ITargetPreparer
+        // preparer, it is not considered as a IBuildReceiver instance.
+        mBuildInfo = testInfo.getBuildInfo();
+
+        // Executes the clean up step.
+        cleanUpModules();
+    }
+
+    private void generateModules() throws IOException {
+        String templateContent = mResourceLoader.load(mTemplate);
+
+        for (String packageName : mPackages) {
+            validatePackageName(packageName);
+            Files.write(
+                    getModulePath(packageName),
+                    templateContent.replaceAll(TEMPLATE_PACKAGE_PATTERN, packageName).getBytes());
+        }
+    }
+
+    private void cleanUpModules() {
+        mPackages.forEach(
+                packageName -> {
+                    try {
+                        Files.delete(getModulePath(packageName));
+                    } catch (IOException ioException) {
+                        CLog.e(
+                                "Failed to delete the generated module for package " + packageName,
+                                ioException);
+                    }
+                });
+    }
+
+    private Path getModulePath(String packageName) throws IOException {
+        Path testsDir = mTestDirectoryProvider.get(mBuildInfo);
+        return testsDir.resolve(packageName + MODULE_FILE_EXTENSION);
+    }
+
+    private static void validatePackageName(String packageName) {
+        if (packageName.isEmpty() || packageName.matches(".*" + TEMPLATE_PACKAGE_PATTERN + ".*")) {
+            throw new IllegalArgumentException(
+                    "Package name cannot be empty or contains package placeholder: "
+                            + TEMPLATE_PACKAGE_PATTERN);
+        }
+    }
+
+    @VisibleForTesting
+    interface ResourceLoader {
+        String load(String resourceName) throws IOException;
+    }
+
+    private static final class ClassResourceLoader implements ResourceLoader {
+        @Override
+        public String load(String resourceName) throws IOException {
+            return Resources.toString(
+                    getClass().getClassLoader().getResource(resourceName), StandardCharsets.UTF_8);
+        }
+    }
+
+    @VisibleForTesting
+    interface TestDirectoryProvider {
+        Path get(IBuildInfo buildInfo) throws IOException;
+    }
+
+    private static final class CompatibilityTestDirectoryProvider implements TestDirectoryProvider {
+        private final FileSystem mFileSystem;
+
+        private CompatibilityTestDirectoryProvider(FileSystem fileSystem) {
+            mFileSystem = fileSystem;
+        }
+
+        @Override
+        public Path get(IBuildInfo buildInfo) throws IOException {
+            return mFileSystem.getPath(
+                    new CompatibilityBuildHelper(buildInfo).getTestsDir().getPath());
+        }
+    }
+}
diff --git a/harness/src/main/java/com/android/csuite/core/GenerateModulePreparer.java b/harness/src/main/java/com/android/csuite/core/GenerateModulePreparer.java
deleted file mode 100644
index 33d1b85..0000000
--- a/harness/src/main/java/com/android/csuite/core/GenerateModulePreparer.java
+++ /dev/null
@@ -1,166 +0,0 @@
-/*
- * Copyright (C) 2020 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.csuite.core;
-
-import com.android.compatibility.common.tradefed.build.CompatibilityBuildHelper;
-import com.android.tradefed.config.Option;
-import com.android.tradefed.config.Option.Importance;
-import com.android.tradefed.device.DeviceNotAvailableException;
-import com.android.tradefed.invoker.TestInformation;
-import com.android.tradefed.log.LogUtil.CLog;
-import com.android.tradefed.targetprep.ITargetPreparer;
-import com.android.tradefed.targetprep.TargetSetupError;
-
-import com.google.common.annotations.VisibleForTesting;
-import com.google.common.io.Resources;
-
-import java.io.IOException;
-import java.nio.charset.StandardCharsets;
-import java.nio.file.FileSystem;
-import java.nio.file.FileSystems;
-import java.nio.file.Files;
-import java.nio.file.Path;
-import java.util.ArrayList;
-import java.util.HashSet;
-import java.util.List;
-import java.util.Set;
-
-/**
- * A preparer for generating TradeFed test suite modules config files.
- *
- * <p>This preparer generates module config files into TradeFed's test directory at runtime using a
- * template. The entire test directory is delete before every run. As a result, there can only be
- * one instance executing at a given time.
- */
-public final class GenerateModulePreparer implements ITargetPreparer {
-
-    @VisibleForTesting static final String MODULE_FILE_EXTENSION = ".config";
-    @VisibleForTesting static final String OPTION_TEMPLATE = "template";
-    @VisibleForTesting static final String OPTION_PACKAGE = "package";
-    private static final String TEMPLATE_PACKAGE_PATTERN = "\\{package\\}";
-
-    @Option(
-            name = OPTION_TEMPLATE,
-            description = "Module config template resource path.",
-            importance = Importance.ALWAYS)
-    private String mTemplate;
-
-    @Option(name = OPTION_PACKAGE, description = "App package names.")
-    private final Set<String> mPackages = new HashSet<>();
-
-    private final TestDirectoryProvider mTestDirectoryProvider;
-    private final ResourceLoader mResourceLoader;
-    private final FileSystem mFileSystem;
-    private final List<Path> mGeneratedModules = new ArrayList<>();
-
-    public GenerateModulePreparer() {
-        this(FileSystems.getDefault());
-    }
-
-    private GenerateModulePreparer(FileSystem fileSystem) {
-        this(
-                fileSystem,
-                new CompatibilityTestDirectoryProvider(fileSystem),
-                new ClassResourceLoader());
-    }
-
-    @VisibleForTesting
-    GenerateModulePreparer(
-            FileSystem fileSystem,
-            TestDirectoryProvider testDirectoryProvider,
-            ResourceLoader resourceLoader) {
-        mFileSystem = fileSystem;
-        mTestDirectoryProvider = testDirectoryProvider;
-        mResourceLoader = resourceLoader;
-    }
-
-    /** {@inheritDoc} */
-    @Override
-    public void setUp(TestInformation testInfo)
-            throws TargetSetupError, DeviceNotAvailableException {
-        try {
-            Path testsDir = mTestDirectoryProvider.get(testInfo);
-            String templateContent = mResourceLoader.load(mTemplate);
-
-            for (String packageName : mPackages) {
-                validatePackageName(packageName);
-                Path modulePath = testsDir.resolve(packageName + MODULE_FILE_EXTENSION);
-                Files.write(
-                        modulePath,
-                        templateContent
-                                .replaceAll(TEMPLATE_PACKAGE_PATTERN, packageName)
-                                .getBytes());
-                mGeneratedModules.add(modulePath);
-            }
-        } catch (IOException e) {
-            throw new TargetSetupError("Failed to generate modules", e);
-        }
-    }
-
-    /** {@inheritDoc} */
-    @Override
-    public void tearDown(TestInformation testInfo, Throwable e) throws DeviceNotAvailableException {
-        mGeneratedModules.forEach(
-                modulePath -> {
-                    try {
-                        Files.delete(modulePath);
-                    } catch (IOException ioException) {
-                        CLog.e("Failed to delete a generated module: " + modulePath, ioException);
-                    }
-                });
-    }
-
-    private static void validatePackageName(String packageName) throws TargetSetupError {
-        if (packageName.isEmpty() || packageName.matches(".*" + TEMPLATE_PACKAGE_PATTERN + ".*")) {
-            throw new TargetSetupError(
-                    "Package name cannot be empty or contains package placeholder: "
-                            + TEMPLATE_PACKAGE_PATTERN);
-        }
-    }
-
-    @VisibleForTesting
-    interface ResourceLoader {
-        String load(String resourceName) throws IOException;
-    }
-
-    private static final class ClassResourceLoader implements ResourceLoader {
-        @Override
-        public String load(String resourceName) throws IOException {
-            return Resources.toString(
-                    getClass().getClassLoader().getResource(resourceName), StandardCharsets.UTF_8);
-        }
-    }
-
-    @VisibleForTesting
-    interface TestDirectoryProvider {
-        Path get(TestInformation testInfo) throws IOException;
-    }
-
-    private static final class CompatibilityTestDirectoryProvider implements TestDirectoryProvider {
-        private final FileSystem mFileSystem;
-
-        private CompatibilityTestDirectoryProvider(FileSystem fileSystem) {
-            mFileSystem = fileSystem;
-        }
-
-        @Override
-        public Path get(TestInformation testInfo) throws IOException {
-            return mFileSystem.getPath(
-                    new CompatibilityBuildHelper(testInfo.getBuildInfo()).getTestsDir().getPath());
-        }
-    }
-}
diff --git a/harness/src/main/resources/config/csuite-app-launch.xml b/harness/src/main/resources/config/csuite-app-launch.xml
index dacfe87..979fc3c 100644
--- a/harness/src/main/resources/config/csuite-app-launch.xml
+++ b/harness/src/main/resources/config/csuite-app-launch.xml
@@ -14,9 +14,10 @@
      limitations under the License.
 -->
 <configuration description="C-Suite Compatibility Launch Test Plan">
+  <test class="com.android.csuite.config.ModuleGenerator"  >
+    <option name="template" value="config/csuite-app-launch.xml.template" />
+  </test>
+  <target_preparer class="com.android.csuite.config.ModuleGenerator" />
   <include name="csuite-base" />
   <option name="plan" value="csuite-app-launch" />
-  <target_preparer class="com.android.csuite.core.GenerateModulePreparer" >
-    <option name="template" value="config/csuite-app-launch.xml.template" />
-  </target_preparer>
 </configuration>
diff --git a/harness/src/main/resources/config/csuite-system-app-launch.xml b/harness/src/main/resources/config/csuite-system-app-launch.xml
index 5838714..28da404 100644
--- a/harness/src/main/resources/config/csuite-system-app-launch.xml
+++ b/harness/src/main/resources/config/csuite-system-app-launch.xml
@@ -14,9 +14,10 @@
      limitations under the License.
 -->
 <configuration description="C-Suite Compatibility Launch Test Plan">
+  <test class="com.android.csuite.config.ModuleGenerator"  >
+    <option name="template" value="config/csuite-system-app-launch.xml.template" />
+  </test>
+  <target_preparer class="com.android.csuite.config.ModuleGenerator" />
   <include name="csuite-base" />
   <option name="plan" value="csuite-system-app-launch" />
-  <target_preparer class="com.android.csuite.core.GenerateModulePreparer" >
-    <option name="template" value="config/csuite-system-app-launch.xml.template" />
-  </target_preparer>
 </configuration>
\ No newline at end of file
diff --git a/harness/src/main/resources/config/csuite-test-package-launch.xml b/harness/src/main/resources/config/csuite-test-package-launch.xml
index f52ba3f..c7b1902 100644
--- a/harness/src/main/resources/config/csuite-test-package-launch.xml
+++ b/harness/src/main/resources/config/csuite-test-package-launch.xml
@@ -14,9 +14,10 @@
      limitations under the License.
 -->
 <configuration description="C-Suite Compatibility Launch Test Plan">
+  <test class="com.android.csuite.config.ModuleGenerator"  >
+    <option name="template" value="config/csuite-test-package-launch.xml.template" />
+  </test>
+  <target_preparer class="com.android.csuite.config.ModuleGenerator" />
   <include name="csuite-base" />
   <option name="plan" value="csuite-test-package-launch" />
-  <target_preparer class="com.android.csuite.core.GenerateModulePreparer" >
-    <option name="template" value="config/csuite-test-package-launch.xml.template" />
-  </target_preparer>
 </configuration>
\ No newline at end of file
diff --git a/harness/src/test/java/com/android/csuite/CSuiteUnitTests.java b/harness/src/test/java/com/android/csuite/CSuiteUnitTests.java
index 01df8a6..096f882 100644
--- a/harness/src/test/java/com/android/csuite/CSuiteUnitTests.java
+++ b/harness/src/test/java/com/android/csuite/CSuiteUnitTests.java
@@ -26,7 +26,7 @@
     com.android.compatibility.targetprep.SystemAppRemovalPreparerTest.class,
     com.android.compatibility.testtype.AppLaunchTestTest.class,
     com.android.csuite.config.AppRemoteFileResolverTest.class,
-    com.android.csuite.core.GenerateModulePreparerTest.class,
+    com.android.csuite.config.ModuleGeneratorTest.class,
     com.android.csuite.testing.CorrespondencesTest.class,
     com.android.csuite.testing.MoreAssertsTest.class,
 })
diff --git a/harness/src/test/java/com/android/csuite/core/GenerateModulePreparerTest.java b/harness/src/test/java/com/android/csuite/config/ModuleGeneratorTest.java
similarity index 66%
rename from harness/src/test/java/com/android/csuite/core/GenerateModulePreparerTest.java
rename to harness/src/test/java/com/android/csuite/config/ModuleGeneratorTest.java
index eeaa2e2..6834bd8 100644
--- a/harness/src/test/java/com/android/csuite/core/GenerateModulePreparerTest.java
+++ b/harness/src/test/java/com/android/csuite/config/ModuleGeneratorTest.java
@@ -13,15 +13,19 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package com.android.csuite.core;
+
+package com.android.csuite.config;
 
 import static com.google.common.truth.Truth.assertThat;
 
 import static org.testng.Assert.assertThrows;
 
+import com.android.tradefed.build.BuildInfo;
 import com.android.tradefed.config.OptionSetter;
+import com.android.tradefed.device.ITestDevice;
+import com.android.tradefed.invoker.IInvocationContext;
+import com.android.tradefed.invoker.InvocationContext;
 import com.android.tradefed.invoker.TestInformation;
-import com.android.tradefed.targetprep.TargetSetupError;
 
 import com.google.common.collect.ArrayListMultimap;
 import com.google.common.collect.ImmutableList;
@@ -34,6 +38,7 @@
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
+import org.mockito.Mockito;
 
 import java.io.IOException;
 import java.nio.file.FileSystem;
@@ -44,7 +49,7 @@
 import java.util.Map;
 
 @RunWith(JUnit4.class)
-public final class GenerateModulePreparerTest {
+public final class ModuleGeneratorTest {
     private static final String TEST_PACKAGE_NAME1 = "test.package.name1";
     private static final String TEST_PACKAGE_NAME2 = "test.package.name2";
     private static final String PACKAGE_PLACEHOLDER = "{package}";
@@ -54,59 +59,63 @@
 
     @Test
     public void tearDown_packageOptionIsSet_deletesGeneratedModules() throws Exception {
-        TestInformation testInfo = createTestInfo();
         Path testsDir = createTestsDir();
-        GenerateModulePreparer preparer =
-                createPreparerBuilder()
+        ModuleGenerator generator1 =
+                createGeneratorBuilder()
                         .setTestsDir(testsDir)
                         .addPackage(TEST_PACKAGE_NAME1)
                         .addPackage(TEST_PACKAGE_NAME1) // Simulate duplicate package option
                         .addPackage(TEST_PACKAGE_NAME2)
                         .build();
-        preparer.setUp(testInfo);
+        generator1.split();
+        ModuleGenerator generator2 =
+                createGeneratorBuilder()
+                        .setTestsDir(testsDir)
+                        .addPackage(TEST_PACKAGE_NAME1)
+                        .addPackage(TEST_PACKAGE_NAME1) // Simulate duplicate package option
+                        .addPackage(TEST_PACKAGE_NAME2)
+                        .build();
 
-        preparer.tearDown(testInfo, NO_EXCEPTION);
+        generator2.tearDown(createTestInfo(), NO_EXCEPTION);
 
         assertThatListDirectory(testsDir).isEmpty();
     }
 
     @Test
     public void tearDown_packageOptionIsNotSet_doesNotThrowError() throws Exception {
-        TestInformation testInfo = createTestInfo();
-        GenerateModulePreparer preparer =
-                createPreparerBuilder().setTestsDir(createTestsDir()).build();
-        preparer.setUp(testInfo);
+        ModuleGenerator generator = createGeneratorBuilder().setTestsDir(createTestsDir()).build();
+        generator.split();
 
-        preparer.tearDown(testInfo, NO_EXCEPTION);
+        generator.tearDown(createTestInfo(), NO_EXCEPTION);
     }
 
     @Test
-    public void setUp_packageNameIsEmptyString_throwsError() throws Exception {
-        GenerateModulePreparer preparer = createPreparerBuilder().addPackage("").build();
+    public void split_packageNameIsEmptyString_throwsError() throws Exception {
+        ModuleGenerator generator = createGeneratorBuilder().addPackage("").build();
 
-        assertThrows(TargetSetupError.class, () -> preparer.setUp(createTestInfo()));
+        assertThrows(IllegalArgumentException.class, () -> generator.split());
     }
 
     @Test
-    public void setUp_packageNameContainsPlaceholder_throwsError() throws Exception {
-        GenerateModulePreparer preparer =
-                createPreparerBuilder().addPackage("a" + PACKAGE_PLACEHOLDER + "b").build();
+    public void split_packageNameContainsPlaceholder_throwsError() throws Exception {
+        ModuleGenerator generator =
+                createGeneratorBuilder().addPackage("a" + PACKAGE_PLACEHOLDER + "b").build();
 
-        assertThrows(TargetSetupError.class, () -> preparer.setUp(createTestInfo()));
+        assertThrows(IllegalArgumentException.class, () -> generator.split());
     }
 
     @Test
-    public void setUp_packageOptionContainsDuplicates_ignoreDuplicates() throws Exception {
+    public void split_packageOptionContainsDuplicates_ignoreDuplicates() throws Exception {
         Path testsDir = createTestsDir();
-        GenerateModulePreparer preparer =
-                createPreparerBuilder()
+        ModuleGenerator generator =
+                createGeneratorBuilder()
                         .setTestsDir(testsDir)
                         .addPackage(TEST_PACKAGE_NAME1)
                         .addPackage(TEST_PACKAGE_NAME1) // Simulate duplicate package option
                         .addPackage(TEST_PACKAGE_NAME2)
                         .build();
 
-        preparer.setUp(createTestInfo());
+        generator.split();
 
         assertThatListDirectory(testsDir)
                 .containsExactly(
@@ -115,21 +124,21 @@
     }
 
     @Test
-    public void setUp_packageOptionNotSet_doesNotGenerate() throws Exception {
+    public void split_packageOptionNotSet_doesNotGenerate() throws Exception {
         Path testsDir = createTestsDir();
-        GenerateModulePreparer preparer = createPreparerBuilder().setTestsDir(testsDir).build();
+        ModuleGenerator generator = createGeneratorBuilder().setTestsDir(testsDir).build();
 
-        preparer.setUp(createTestInfo());
+        generator.split();
 
         assertThatListDirectory(testsDir).isEmpty();
     }
 
     @Test
-    public void setUp_templateContainsPlaceholders_replacesPlaceholdersInOutput() throws Exception {
+    public void split_templateContainsPlaceholders_replacesPlaceholdersInOutput() throws Exception {
         Path testsDir = createTestsDir();
         String content = "hello placeholder%s%s world";
-        GenerateModulePreparer preparer =
-                createPreparerBuilder()
+        ModuleGenerator generator =
+                createGeneratorBuilder()
                         .setTestsDir(testsDir)
                         .addPackage(TEST_PACKAGE_NAME1)
                         .addPackage(TEST_PACKAGE_NAME2)
@@ -137,7 +146,7 @@
                                 String.format(content, PACKAGE_PLACEHOLDER, PACKAGE_PLACEHOLDER))
                         .build();
 
-        preparer.setUp(createTestInfo());
+        generator.split();
 
         assertThatModuleConfigFileContent(testsDir, TEST_PACKAGE_NAME1)
                 .isEqualTo(String.format(content, TEST_PACKAGE_NAME1, TEST_PACKAGE_NAME1));
@@ -146,36 +155,36 @@
     }
 
     @Test
-    public void setUp_templateDoesNotContainPlaceholder_outputsTemplateContent() throws Exception {
+    public void split_templateDoesNotContainPlaceholder_outputsTemplateContent() throws Exception {
         Path testsDir = createTestsDir();
         String content = "no placeholder";
-        GenerateModulePreparer preparer =
-                createPreparerBuilder()
+        ModuleGenerator generator =
+                createGeneratorBuilder()
                         .setTestsDir(testsDir)
                         .addPackage(TEST_PACKAGE_NAME1)
                         .addPackage(TEST_PACKAGE_NAME2)
                         .setTemplateContent(content)
                         .build();
 
-        preparer.setUp(createTestInfo());
+        generator.split();
 
         assertThatModuleConfigFileContent(testsDir, TEST_PACKAGE_NAME1).isEqualTo(content);
         assertThatModuleConfigFileContent(testsDir, TEST_PACKAGE_NAME2).isEqualTo(content);
     }
 
     @Test
-    public void setUp_templateContentIsEmpty_outputsTemplateContent() throws Exception {
+    public void split_templateContentIsEmpty_outputsTemplateContent() throws Exception {
         Path testsDir = createTestsDir();
         String content = "";
-        GenerateModulePreparer preparer =
-                createPreparerBuilder()
+        ModuleGenerator generator =
+                createGeneratorBuilder()
                         .setTestsDir(testsDir)
                         .addPackage(TEST_PACKAGE_NAME1)
                         .addPackage(TEST_PACKAGE_NAME2)
                         .setTemplateContent(content)
                         .build();
 
-        preparer.setUp(createTestInfo());
+        generator.split();
 
         assertThatModuleConfigFileContent(testsDir, TEST_PACKAGE_NAME1).isEqualTo(content);
         assertThatModuleConfigFileContent(testsDir, TEST_PACKAGE_NAME2).isEqualTo(content);
@@ -206,63 +215,66 @@
     }
 
     private static TestInformation createTestInfo() {
-        return TestInformation.newBuilder().build();
+        IInvocationContext context = new InvocationContext();
+        context.addAllocatedDevice("device1", Mockito.mock(ITestDevice.class));
+        context.addDeviceBuildInfo("device1", new BuildInfo());
+        return TestInformation.newBuilder().setInvocationContext(context).build();
     }
 
-    private PreparerBuilder createPreparerBuilder() throws IOException {
-        return new PreparerBuilder()
+    private GeneratorBuilder createGeneratorBuilder() throws IOException {
+        return new GeneratorBuilder()
                 .setFileSystem(mFileSystem)
                 .setTemplateContent(MODULE_TEMPLATE_CONTENT)
-                .setOption(GenerateModulePreparer.OPTION_TEMPLATE, "empty_path");
+                .setOption(ModuleGenerator.OPTION_TEMPLATE, "empty_path");
     }
 
-    private static final class PreparerBuilder {
+    private static final class GeneratorBuilder {
         private final ListMultimap<String, String> mOptions = ArrayListMultimap.create();
         private final List<String> mPackages = new ArrayList<>();
         private Path mTestsDir;
         private String mTemplateContent;
         private FileSystem mFileSystem;
 
-        PreparerBuilder addPackage(String packageName) {
+        GeneratorBuilder addPackage(String packageName) {
             mPackages.add(packageName);
             return this;
         }
 
-        PreparerBuilder setFileSystem(FileSystem fileSystem) {
+        GeneratorBuilder setFileSystem(FileSystem fileSystem) {
             mFileSystem = fileSystem;
             return this;
         }
 
-        PreparerBuilder setTemplateContent(String templateContent) {
+        GeneratorBuilder setTemplateContent(String templateContent) {
             mTemplateContent = templateContent;
             return this;
         }
 
-        PreparerBuilder setTestsDir(Path testsDir) {
+        GeneratorBuilder setTestsDir(Path testsDir) {
             mTestsDir = testsDir;
             return this;
         }
 
-        PreparerBuilder setOption(String key, String value) {
+        GeneratorBuilder setOption(String key, String value) {
             mOptions.put(key, value);
             return this;
         }
 
-        GenerateModulePreparer build() throws Exception {
-            GenerateModulePreparer preparer =
-                    new GenerateModulePreparer(
-                            mFileSystem, testInfo -> mTestsDir, resourcePath -> mTemplateContent);
+        ModuleGenerator build() throws Exception {
+            ModuleGenerator generator =
+                    new ModuleGenerator(
+                            mFileSystem, buildInfo -> mTestsDir, resourcePath -> mTemplateContent);
 
-            OptionSetter optionSetter = new OptionSetter(preparer);
+            OptionSetter optionSetter = new OptionSetter(generator);
             for (Map.Entry<String, String> entry : mOptions.entries()) {
                 optionSetter.setOptionValue(entry.getKey(), entry.getValue());
             }
 
             for (String packageName : mPackages) {
-                optionSetter.setOptionValue(GenerateModulePreparer.OPTION_PACKAGE, packageName);
+                optionSetter.setOptionValue(ModuleGenerator.OPTION_PACKAGE, packageName);
             }
 
-            return preparer;
+            return generator;
         }
     }
 
@@ -270,9 +282,9 @@
             "<?xml version=\"1.0\" encoding=\"utf-8\"?>\n"
                     + "<configuration description=\"description\">\n"
                     + "    <option name=\"package-name\" value=\"{package}\"/>\n"
-                    + "    <target_preparer class=\"some.preparer.class\">\n"
+                    + "    <target_generator class=\"some.generator.class\">\n"
                     + "        <option name=\"test-file-name\" value=\"app://{package}\"/>\n"
-                    + "    </target_preparer>\n"
+                    + "    </target_generator>\n"
                     + "    <test class=\"some.test.class\"/>\n"
                     + "</configuration>";
 }