Internal refactor

RELNOTES=N/A
PiperOrigin-RevId: 355752966

2 files changed, 23 insertions, 22 deletions

diff --git a/java/dagger/hilt/processor/internal/root/TestComponentDataGenerator.java b/java/dagger/hilt/processor/internal/root/TestComponentDataGenerator.java
index c9e8cfb5b..284f8cdee 100644
--- a/java/dagger/hilt/processor/internal/root/TestComponentDataGenerator.java
+++ b/java/dagger/hilt/processor/internal/root/TestComponentDataGenerator.java
@@ -30,7 +30,6 @@ import com.squareup.javapoet.ClassName;
 import com.squareup.javapoet.CodeBlock;
 import com.squareup.javapoet.JavaFile;
 import com.squareup.javapoet.MethodSpec;
-import com.squareup.javapoet.ParameterizedTypeName;
 import com.squareup.javapoet.TypeSpec;
 import dagger.hilt.processor.internal.ClassNames;
 import dagger.hilt.processor.internal.ComponentNames;
@@ -196,23 +195,22 @@ public final class TestComponentDataGenerator {
   private MethodSpec getTestInjectInternalMethod() {
     TypeElement testElement = rootMetadata.testRootMetadata().testElement();
     ClassName testName = ClassName.get(testElement);
-    MethodSpec.Builder builder =
-        MethodSpec.methodBuilder("injectInternal")
-            .addModifiers(PRIVATE, STATIC)
-            .addParameter(testName, "testInstance")
-            .addAnnotation(
-                AnnotationSpec.builder(SuppressWarnings.class)
-                    .addMember("value", "$S", "unchecked")
-                    .build());
-
-    return builder
-        .addStatement(
-            "(($T) (($T) $T.getApplicationContext()).generatedComponent())"
-                + ".injectTest(testInstance)",
-            ParameterizedTypeName.get(
-                ClassNames.TEST_INJECTOR, rootMetadata.testRootMetadata().testName()),
-            ClassNames.GENERATED_COMPONENT_MANAGER,
-            ClassNames.APPLICATION_PROVIDER)
+    return MethodSpec.methodBuilder("injectInternal")
+        .addModifiers(PRIVATE, STATIC)
+        .addParameter(testName, "testInstance")
+        .addAnnotation(
+            AnnotationSpec.builder(SuppressWarnings.class)
+                .addMember("value", "$S", "unchecked")
+                .build())
+        .addStatement("$L.injectTest(testInstance)", getInjector(testElement))
         .build();
   }
+
+  private static CodeBlock getInjector(TypeElement testElement) {
+    return CodeBlock.of(
+        "(($T) (($T) $T.getApplicationContext()).generatedComponent())",
+        ClassNames.TEST_INJECTOR,
+        ClassNames.GENERATED_COMPONENT_MANAGER,
+        ClassNames.APPLICATION_PROVIDER);
+  }
 }
diff --git a/java/dagger/hilt/processor/internal/root/TestInjectorGenerator.java b/java/dagger/hilt/processor/internal/root/TestInjectorGenerator.java
index 90c190543..fc97feda6 100644
--- a/java/dagger/hilt/processor/internal/root/TestInjectorGenerator.java
+++ b/java/dagger/hilt/processor/internal/root/TestInjectorGenerator.java
@@ -17,6 +17,7 @@
 package dagger.hilt.processor.internal.root;
 
 import com.squareup.javapoet.AnnotationSpec;
+import com.squareup.javapoet.ClassName;
 import com.squareup.javapoet.JavaFile;
 import com.squareup.javapoet.MethodSpec;
 import com.squareup.javapoet.ParameterizedTypeName;
@@ -26,6 +27,7 @@ import dagger.hilt.processor.internal.Processors;
 import java.io.IOException;
 import javax.annotation.processing.ProcessingEnvironment;
 import javax.lang.model.element.Modifier;
+import javax.lang.model.element.TypeElement;
 
 /** Generates an entry point for a test. */
 public final class TestInjectorGenerator {
@@ -48,10 +50,7 @@ public final class TestInjectorGenerator {
             .addAnnotation(ClassNames.GENERATED_ENTRY_POINT)
             .addAnnotation(
                 AnnotationSpec.builder(ClassNames.INSTALL_IN)
-                    .addMember(
-                        "value",
-                        "$T.class",
-                    ClassNames.SINGLETON_COMPONENT)
+                    .addMember("value", "$T.class", installInComponent(metadata.testElement()))
                     .build())
             .addModifiers(Modifier.PUBLIC)
             .addSuperinterface(
@@ -71,4 +70,8 @@ public final class TestInjectorGenerator {
         .build()
         .writeTo(env.getFiler());
   }
+
+  private static ClassName installInComponent(TypeElement testElement) {
+    return ClassNames.SINGLETON_COMPONENT;
+  }
 }