diff --git a/java/com/android/dialer/inject/DialerRootComponent.java b/java/com/android/dialer/inject/DialerRootComponent.java
index 11dbf96..e508e17 100644
--- a/java/com/android/dialer/inject/DialerRootComponent.java
+++ b/java/com/android/dialer/inject/DialerRootComponent.java
@@ -41,4 +41,6 @@
 @Retention(RetentionPolicy.SOURCE)
 public @interface DialerRootComponent {
   DialerVariant variant();
+
+  Class<?> injectClass() default Object.class;
 }
diff --git a/java/com/android/dialer/rootcomponentgenerator/RootComponentGeneratingStep.java b/java/com/android/dialer/rootcomponentgenerator/RootComponentGeneratingStep.java
index 5321849..06b994f 100644
--- a/java/com/android/dialer/rootcomponentgenerator/RootComponentGeneratingStep.java
+++ b/java/com/android/dialer/rootcomponentgenerator/RootComponentGeneratingStep.java
@@ -26,7 +26,6 @@
 import com.android.dialer.inject.IncludeInDialerRoot;
 import com.android.dialer.inject.InstallIn;
 import com.android.dialer.inject.RootComponentGeneratorMetadata;
-import com.android.dialer.inject.testing.GenerateTestDaggerApp;
 import com.google.auto.common.BasicAnnotationProcessor.ProcessingStep;
 import com.google.auto.common.MoreElements;
 import com.google.common.base.Optional;
@@ -81,27 +80,18 @@
 
   private final ProcessingEnvironment processingEnv;
 
-  private TypeElement annotatedTest;
-
   public RootComponentGeneratingStep(ProcessingEnvironment processingEnv) {
     this.processingEnv = processingEnv;
   }
 
   @Override
   public Set<? extends Class<? extends Annotation>> annotations() {
-    return ImmutableSet.of(
-        DialerRootComponent.class,
-        InstallIn.class,
-        IncludeInDialerRoot.class,
-        GenerateTestDaggerApp.class);
+    return ImmutableSet.of(DialerRootComponent.class, InstallIn.class, IncludeInDialerRoot.class);
   }
 
   @Override
   public Set<? extends Element> process(
       SetMultimap<Class<? extends Annotation>, Element> elementsByAnnotation) {
-    for (Element element : elementsByAnnotation.get(GenerateTestDaggerApp.class)) {
-      annotatedTest = MoreElements.asType(element);
-    }
     for (Element element : elementsByAnnotation.get(DialerRootComponent.class)) {
       // defer root components to the next round in case where the current build target contains
       // elements annotated with @InstallIn. Annotation processor cannot detect metadata files
@@ -144,9 +134,15 @@
       componentAnnotation.addMember("modules", "$T.class", annotatedElement.asType());
     }
     rootComponentClassBuilder.addAnnotation(componentAnnotation.build());
-    if (annotatedTest != null) {
-      rootComponentClassBuilder.addMethod(generateInjectMethod());
-    }
+
+    AnnotationMirror dialerRootComponentMirror =
+        getAnnotationMirror(rootElement, DialerRootComponent.class).get();
+
+    TypeMirror annotatedTestClass =
+        (TypeMirror) getAnnotationValue(dialerRootComponentMirror, "injectClass").getValue();
+
+    rootComponentClassBuilder.addMethod(generateInjectMethod(annotatedTestClass));
+
     TypeSpec rootComponentClass = rootComponentClassBuilder.build();
     RootComponentUtils.writeJavaFile(
         processingEnv, ClassName.get(rootElement).packageName(), rootComponentClass);
@@ -200,11 +196,11 @@
     }
   }
 
-  private MethodSpec generateInjectMethod() {
+  private MethodSpec generateInjectMethod(TypeMirror testClassTypeMirror) {
     return MethodSpec.methodBuilder("inject")
         .addModifiers(Modifier.PUBLIC, Modifier.ABSTRACT)
         .returns(void.class)
-        .addParameter(ClassName.get(annotatedTest), "test")
+        .addParameter(ClassName.get(testClassTypeMirror), "clazz")
         .build();
   }