Merge "Improve HiltViewModelProcessor" into androidx-master-dev

12 files changed, 911 insertions, 47 deletions

diff --git a/lifecycle/lifecycle-viewmodel-hilt-compiler/build.gradle b/lifecycle/lifecycle-viewmodel-hilt-compiler/build.gradle
index bddf73594ed..31328c2302d 100644
--- a/lifecycle/lifecycle-viewmodel-hilt-compiler/build.gradle
+++ b/lifecycle/lifecycle-viewmodel-hilt-compiler/build.gradle
@@ -43,8 +43,14 @@ dependencies {
     testImplementation(JUNIT)
     testImplementation(TRUTH)
     testImplementation(GOOGLE_COMPILE_TESTING)
+    testImplementation fileTree(
+            dir: "${new File(project(":lifecycle:lifecycle-viewmodel-hilt").buildDir, "libJar")}",
+            include : "*.jar")
+    testImplementation(HILT_ANDROID)
 }
 
+tasks.findByName("compileKotlin").dependsOn(":lifecycle:lifecycle-viewmodel-hilt:jarDebug")
+
 androidx {
     name = "Android Lifecycle ViewModel Hilt Extension Compiler"
     publish = Publish.NONE
diff --git a/lifecycle/lifecycle-viewmodel-hilt-compiler/src/main/kotlin/androidx/lifecycle/hilt/HiltViewModelElements.kt b/lifecycle/lifecycle-viewmodel-hilt-compiler/src/main/kotlin/androidx/lifecycle/hilt/HiltViewModelElements.kt
index 3adfd45e47d..c7c9c7cad04 100644
--- a/lifecycle/lifecycle-viewmodel-hilt-compiler/src/main/kotlin/androidx/lifecycle/hilt/HiltViewModelElements.kt
+++ b/lifecycle/lifecycle-viewmodel-hilt-compiler/src/main/kotlin/androidx/lifecycle/hilt/HiltViewModelElements.kt
@@ -16,11 +16,15 @@
 
 package androidx.lifecycle.hilt
 
+import androidx.lifecycle.hilt.ext.hasAnnotation
 import com.google.auto.common.MoreElements
+import com.squareup.javapoet.AnnotationSpec
 import com.squareup.javapoet.ClassName
 import com.squareup.javapoet.ParameterizedTypeName
+import com.squareup.javapoet.TypeName
 import javax.lang.model.element.ExecutableElement
 import javax.lang.model.element.TypeElement
+import javax.lang.model.element.VariableElement
 
 /**
  * Data class that represents a Hilt injected ViewModel
@@ -31,17 +35,49 @@ internal data class HiltViewModelElements(
 ) {
     val className = ClassName.get(typeElement)
 
-    // TODO(danysantiago): Handle nested classes
     val factoryClassName = ClassName.get(
         MoreElements.getPackage(typeElement).qualifiedName.toString(),
-        "${typeElement.simpleName}_AssistedFactory")
+        "${className.simpleNames().joinToString("_")}_AssistedFactory")
 
     val factorySuperTypeName = ParameterizedTypeName.get(
         ClassNames.VIEW_MODEL_ASSISTED_FACTORY,
         className)
 
-    // TODO(danysantiago): Handle nested classes
     val moduleClassName = ClassName.get(
         MoreElements.getPackage(typeElement).qualifiedName.toString(),
-        "${typeElement.simpleName}_HiltModule")
-}
\ No newline at end of file
+        "${className.simpleNames().joinToString("_")}_HiltModule")
+
+    val dependencyRequests = constructorElement.parameters.map { it.toDependencyRequest() }
+}
+
+/**
+ * Data class that represents a binding request from the injected ViewModel
+ */
+internal data class DependencyRequest(
+    val name: String,
+    val type: TypeName,
+    val qualifier: AnnotationSpec? = null
+) {
+    val isProvider = type is ParameterizedTypeName && type.rawType == ClassNames.PROVIDER
+
+    val providerTypeName: TypeName = let {
+        val type = if (isProvider) {
+            type // Do not wrap a Provider inside another Provider.
+        } else {
+            ParameterizedTypeName.get(ClassNames.PROVIDER, type.box())
+        }
+        if (qualifier != null) {
+            type.annotated(qualifier)
+        } else {
+            type
+        }
+    }
+}
+
+internal fun VariableElement.toDependencyRequest() = DependencyRequest(
+    name = simpleName.toString(),
+    type = TypeName.get(asType()),
+    qualifier = annotationMirrors.find {
+        it.annotationType.asElement().hasAnnotation("javax.inject.Qualifier")
+    }?.let { AnnotationSpec.get(it) }
+)
\ No newline at end of file
diff --git a/lifecycle/lifecycle-viewmodel-hilt-compiler/src/main/kotlin/androidx/lifecycle/hilt/HiltViewModelGenerator.kt b/lifecycle/lifecycle-viewmodel-hilt-compiler/src/main/kotlin/androidx/lifecycle/hilt/HiltViewModelGenerator.kt
index aa6a85e3b0c..6961686982c 100644
--- a/lifecycle/lifecycle-viewmodel-hilt-compiler/src/main/kotlin/androidx/lifecycle/hilt/HiltViewModelGenerator.kt
+++ b/lifecycle/lifecycle-viewmodel-hilt-compiler/src/main/kotlin/androidx/lifecycle/hilt/HiltViewModelGenerator.kt
@@ -16,6 +16,10 @@
 
 package androidx.lifecycle.hilt
 
+import androidx.lifecycle.hilt.ext.L
+import androidx.lifecycle.hilt.ext.T
+import androidx.lifecycle.hilt.ext.W
+import androidx.lifecycle.hilt.ext.addGeneratedAnnotation
 import com.squareup.javapoet.AnnotationSpec
 import com.squareup.javapoet.CodeBlock
 import com.squareup.javapoet.FieldSpec
@@ -23,7 +27,6 @@ import com.squareup.javapoet.JavaFile
 import com.squareup.javapoet.MethodSpec
 import com.squareup.javapoet.ParameterSpec
 import com.squareup.javapoet.ParameterizedTypeName
-import com.squareup.javapoet.TypeName
 import com.squareup.javapoet.TypeSpec
 import com.squareup.javapoet.WildcardTypeName
 import javax.annotation.processing.ProcessingEnvironment
@@ -40,7 +43,7 @@ import javax.lang.model.element.Modifier
  *   @Binds
  *   @IntoMap
  *   @ViewModelKey($.class)
- *   ViewModelAssistedFactory<?> bind($_AssistedFactory f)
+ *   ViewModelAssistedFactory<? extends ViewModel> bind($_AssistedFactory factory)
  * }
  * ```
  * and
@@ -58,7 +61,7 @@ import javax.lang.model.element.Modifier
  *     ...
  *   }
  *
- *   @Overrides
+ *   @Override
  *   @NonNull
  *   public $ create(@NonNull SavedStateHandle handle) {
  *     return new $(dep1.get(), dep2.get(), ..., handle);
@@ -71,17 +74,14 @@ internal class HiltViewModelGenerator(
     private val viewModelElements: HiltViewModelElements
 ) {
     fun generate() {
-        val fieldsSpecs = getFieldSpecs(viewModelElements)
-        val constructorSpec = getConstructorMethodSpec(fieldsSpecs)
-        val createMethodSpec = getCreateMethodSpec(viewModelElements)
         val factoryTypeSpec = TypeSpec.classBuilder(viewModelElements.factoryClassName)
             .addOriginatingElement(viewModelElements.typeElement)
             .addSuperinterface(viewModelElements.factorySuperTypeName)
             .addModifiers(Modifier.PUBLIC, Modifier.FINAL)
             .addGeneratedAnnotation(processingEnv.elementUtils, processingEnv.sourceVersion)
-            .addFields(fieldsSpecs)
-            .addMethod(constructorSpec)
-            .addMethod(createMethodSpec)
+            .addFields(getFieldSpecs())
+            .addMethod(getConstructorMethodSpec())
+            .addMethod(getCreateMethodSpec())
             .build()
         JavaFile.builder(viewModelElements.factoryClassName.packageName(), factoryTypeSpec)
             .build()
@@ -95,6 +95,7 @@ internal class HiltViewModelGenerator(
                 AnnotationSpec.builder(ClassNames.INSTALL_IN)
                     .addMember("value", "$T.class", ClassNames.ACTIVITY_RETAINED_COMPONENT)
                     .build())
+            .addModifiers(Modifier.PUBLIC)
             .addMethod(
                 MethodSpec.methodBuilder("bind")
                     .addAnnotation(ClassNames.BINDS)
@@ -116,41 +117,34 @@ internal class HiltViewModelGenerator(
             .writeTo(processingEnv.filer)
     }
 
-    private fun getFieldSpecs(viewModelElements: HiltViewModelElements) =
-        viewModelElements.constructorElement.parameters.mapNotNull { parameter ->
-            val paramTypeName = TypeName.get(parameter.asType())
-            if (paramTypeName == ClassNames.SAVED_STATE_HANDLE) {
-                // Skip SavedStateHandle since it is assisted injected.
-                return@mapNotNull null
-            }
-            // TODO(danysantiago): Handle qualifiers
-            // TODO(danysantiago): Don't wrap params that are already a Provider
-            FieldSpec.builder(
-                ParameterizedTypeName.get(ClassNames.PROVIDER, paramTypeName),
-                "${parameter.simpleName}Provider",
-                Modifier.PRIVATE, Modifier.FINAL)
+    private fun getFieldSpecs() = viewModelElements.dependencyRequests
+        .filterNot { it.isSavedStateHandle }
+        .map { dependencyRequest ->
+            val fieldTypeName = dependencyRequest.providerTypeName.withoutAnnotations()
+            FieldSpec.builder(fieldTypeName, dependencyRequest.name)
+                .addModifiers(Modifier.PRIVATE, Modifier.FINAL)
                 .build()
         }
 
-    private fun getConstructorMethodSpec(fieldsSpecs: List<FieldSpec>) =
+    private fun getConstructorMethodSpec() =
         MethodSpec.constructorBuilder()
             .addAnnotation(ClassNames.INJECT)
             .apply {
-                fieldsSpecs.forEach { field ->
-                    addParameter(field.type, field.name)
-                    addStatement("this.$1N = $1N", field)
+                viewModelElements.dependencyRequests
+                    .filterNot { it.isSavedStateHandle }
+                    .forEach { dependencyRequest ->
+                        addParameter(dependencyRequest.providerTypeName, dependencyRequest.name)
+                        addStatement("this.$1N = $1N", dependencyRequest.name)
                 }
             }
             .build()
 
-    private fun getCreateMethodSpec(viewModelElements: HiltViewModelElements): MethodSpec {
-        val constructorArgs = viewModelElements.constructorElement.parameters.map { param ->
-            val paramTypeName = TypeName.get(param.asType())
-            val paramLiteral = if (paramTypeName == ClassNames.SAVED_STATE_HANDLE) {
-                "handle"
-            } else {
-                // TODO(danysantiago): Consider using the field specs?
-                "${param.simpleName}Provider.get()"
+    private fun getCreateMethodSpec(): MethodSpec {
+        val constructorArgs = viewModelElements.dependencyRequests.map { dependencyRequest ->
+            val paramLiteral = when {
+                dependencyRequest.isSavedStateHandle -> "handle"
+                dependencyRequest.isProvider -> dependencyRequest.name
+                else -> "${dependencyRequest.name}.get()"
             }
             CodeBlock.of(L, paramLiteral)
         }
@@ -167,4 +161,7 @@ internal class HiltViewModelGenerator(
                 viewModelElements.className, CodeBlock.join(constructorArgs, ",$W"))
             .build()
     }
-}
\ No newline at end of file
+}
+
+internal val DependencyRequest.isSavedStateHandle: Boolean
+    get() = type == ClassNames.SAVED_STATE_HANDLE && qualifier == null
diff --git a/lifecycle/lifecycle-viewmodel-hilt-compiler/src/main/kotlin/androidx/lifecycle/hilt/HiltViewModelProcessor.kt b/lifecycle/lifecycle-viewmodel-hilt-compiler/src/main/kotlin/androidx/lifecycle/hilt/HiltViewModelProcessor.kt
index e61373fc69b..c8f7069eb31 100644
--- a/lifecycle/lifecycle-viewmodel-hilt-compiler/src/main/kotlin/androidx/lifecycle/hilt/HiltViewModelProcessor.kt
+++ b/lifecycle/lifecycle-viewmodel-hilt-compiler/src/main/kotlin/androidx/lifecycle/hilt/HiltViewModelProcessor.kt
@@ -16,6 +16,7 @@
 
 package androidx.lifecycle.hilt
 
+import androidx.lifecycle.hilt.ext.hasAnnotation
 import com.google.auto.common.BasicAnnotationProcessor
 import com.google.auto.common.MoreElements
 import com.google.auto.service.AutoService
@@ -27,6 +28,11 @@ import javax.annotation.processing.Processor
 import javax.lang.model.SourceVersion
 import javax.lang.model.element.Element
 import javax.lang.model.element.ExecutableElement
+import javax.lang.model.element.Modifier
+import javax.lang.model.element.NestingKind
+import javax.lang.model.element.TypeElement
+import javax.lang.model.util.ElementFilter
+import javax.tools.Diagnostic
 
 /**
  * Annotation processor that generates code enabling assisted injection of ViewModels using Hilt.
@@ -43,24 +49,67 @@ class ViewModelInjectStep(
     private val processingEnv: ProcessingEnvironment
 ) : BasicAnnotationProcessor.ProcessingStep {
 
+    private val elements = processingEnv.elementUtils
+    private val types = processingEnv.typeUtils
+    private val messager = processingEnv.messager
+
     override fun annotations() = setOf(ViewModelInject::class.java)
 
     override fun process(
         elementsByAnnotation: SetMultimap<Class<out Annotation>, Element>
     ): MutableSet<out Element> {
+        val parsedElements = mutableSetOf<TypeElement>()
         elementsByAnnotation[ViewModelInject::class.java].forEach { element ->
             val constructorElement = MoreElements.asExecutable(element)
-            parse(constructorElement)?.let { viewModel ->
-                HiltViewModelGenerator(processingEnv, viewModel).generate()
+            val typeElement = MoreElements.asType(constructorElement.enclosingElement)
+            if (parsedElements.add(typeElement)) {
+                parse(typeElement, constructorElement)?.let { viewModel ->
+                    HiltViewModelGenerator(processingEnv, viewModel).generate()
+                }
             }
         }
         return mutableSetOf()
     }
 
-    private fun parse(constructorElement: ExecutableElement): HiltViewModelElements? {
-        val typeElement = MoreElements.asType(constructorElement.enclosingElement)
-        // TODO(danysantiago): Validate type extends ViewModel
-        // TODO(danysantiago): Validate only one constructor is annotated
+    private fun parse(
+        typeElement: TypeElement,
+        constructorElement: ExecutableElement
+    ): HiltViewModelElements? {
+        var valid = true
+
+        if (!types.isSubtype(typeElement.asType(),
+                elements.getTypeElement(ClassNames.VIEW_MODEL.toString()).asType())) {
+            error("@ViewModelInject is only supported on types that subclass " +
+                    "androidx.lifecycle.ViewModel.")
+            valid = false
+        }
+
+        ElementFilter.constructorsIn(typeElement.enclosedElements).filter {
+            it.hasAnnotation(ViewModelInject::class)
+        }.let { constructors ->
+            if (constructors.size > 1) {
+                error("Multiple @ViewModelInject annotated constructors found.", typeElement)
+                valid = false
+            }
+            constructors.filter { it.modifiers.contains(Modifier.PRIVATE) }.forEach {
+                error("@ViewModelInject annotated constructors must not be private.", it)
+                valid = false
+            }
+        }
+
+        if (typeElement.nestingKind == NestingKind.MEMBER &&
+            !typeElement.modifiers.contains(Modifier.STATIC)) {
+            error("@ViewModelInject may only be used on inner classes if they are static.",
+                typeElement)
+            valid = false
+        }
+
+        if (!valid) return null
+
         return HiltViewModelElements(typeElement, constructorElement)
     }
+
+    private fun error(message: String, element: Element? = null) {
+        messager.printMessage(Diagnostic.Kind.ERROR, message, element)
+    }
 }
\ No newline at end of file
diff --git a/lifecycle/lifecycle-viewmodel-hilt-compiler/src/main/kotlin/androidx/lifecycle/hilt/ext/annotationProcessing.kt b/lifecycle/lifecycle-viewmodel-hilt-compiler/src/main/kotlin/androidx/lifecycle/hilt/ext/annotationProcessing.kt
new file mode 100644
index 00000000000..cbddb974a69
--- /dev/null
+++ b/lifecycle/lifecycle-viewmodel-hilt-compiler/src/main/kotlin/androidx/lifecycle/hilt/ext/annotationProcessing.kt
@@ -0,0 +1,28 @@
+/*
+ * Copyright 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 androidx.lifecycle.hilt.ext
+
+import com.google.auto.common.MoreElements
+import javax.lang.model.element.Element
+import kotlin.reflect.KClass
+
+fun Element.hasAnnotation(clazz: KClass<out Annotation>) =
+    MoreElements.isAnnotationPresent(this, clazz.java)
+
+fun Element.hasAnnotation(qName: String) = annotationMirrors.any {
+    MoreElements.asType(it.annotationType.asElement()).qualifiedName.contentEquals(qName)
+}
diff --git a/lifecycle/lifecycle-viewmodel-hilt-compiler/src/main/kotlin/androidx/lifecycle/hilt/JavaPoetExt.kt b/lifecycle/lifecycle-viewmodel-hilt-compiler/src/main/kotlin/androidx/lifecycle/hilt/ext/javaPoet.kt
index 669b6b7bfdf..baa0d01b00f 100644
--- a/lifecycle/lifecycle-viewmodel-hilt-compiler/src/main/kotlin/androidx/lifecycle/hilt/JavaPoetExt.kt
+++ b/lifecycle/lifecycle-viewmodel-hilt-compiler/src/main/kotlin/androidx/lifecycle/hilt/ext/javaPoet.kt
@@ -14,8 +14,9 @@
  * limitations under the License.
  */
 
-package androidx.lifecycle.hilt
+package androidx.lifecycle.hilt.ext
 
+import androidx.lifecycle.hilt.HiltViewModelProcessor
 import com.google.auto.common.GeneratedAnnotationSpecs
 import com.squareup.javapoet.TypeSpec
 import javax.lang.model.SourceVersion
diff --git a/lifecycle/lifecycle-viewmodel-hilt-compiler/src/test/data/sources/SavedStateHandle.java b/lifecycle/lifecycle-viewmodel-hilt-compiler/src/test/data/sources/SavedStateHandle.java
new file mode 100644
index 00000000000..72659cc4553
--- /dev/null
+++ b/lifecycle/lifecycle-viewmodel-hilt-compiler/src/test/data/sources/SavedStateHandle.java
@@ -0,0 +1,21 @@
+/*
+ * Copyright 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 androidx.lifecycle;
+
+public class SavedStateHandle {
+
+}
diff --git a/lifecycle/lifecycle-viewmodel-hilt-compiler/src/test/data/sources/ViewModel.java b/lifecycle/lifecycle-viewmodel-hilt-compiler/src/test/data/sources/ViewModel.java
new file mode 100644
index 00000000000..9422d516ea9
--- /dev/null
+++ b/lifecycle/lifecycle-viewmodel-hilt-compiler/src/test/data/sources/ViewModel.java
@@ -0,0 +1,21 @@
+/*
+ * Copyright 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 androidx.lifecycle;
+
+public abstract class ViewModel {
+
+}
diff --git a/lifecycle/lifecycle-viewmodel-hilt-compiler/src/test/kotlin/androidx/lifecycle/hilt/HiltViewModelGeneratorTest.kt b/lifecycle/lifecycle-viewmodel-hilt-compiler/src/test/kotlin/androidx/lifecycle/hilt/HiltViewModelGeneratorTest.kt
new file mode 100644
index 00000000000..851d4837741
--- /dev/null
+++ b/lifecycle/lifecycle-viewmodel-hilt-compiler/src/test/kotlin/androidx/lifecycle/hilt/HiltViewModelGeneratorTest.kt
@@ -0,0 +1,515 @@
+/*
+ * Copyright 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 androidx.lifecycle.hilt
+
+import com.google.testing.compile.CompilationSubject.assertThat
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.junit.runners.JUnit4
+
+@RunWith(JUnit4::class)
+class HiltViewModelGeneratorTest {
+
+    private val GENERATED_TYPE = try {
+        Class.forName("javax.annotation.processing.Generated")
+        "javax.annotation.processing.Generated"
+    } catch (_: ClassNotFoundException) {
+        "javax.annotation.Generated"
+    }
+
+    private val GENERATED_ANNOTATION =
+        "@Generated(\"androidx.lifecycle.hilt.HiltViewModelProcessor\")"
+
+    @Test
+    fun verifyAssistedFactory_noArg() {
+        val myViewModel = """
+        package androidx.lifecycle.hilt.test;
+
+        import androidx.lifecycle.ViewModel;
+        import androidx.lifecycle.hilt.ViewModelInject;
+
+        class MyViewModel extends ViewModel {
+            @ViewModelInject
+            MyViewModel() { }
+        }
+        """.toJFO("androidx.lifecycle.hilt.test.MyViewModel")
+
+        val expected = """
+        package androidx.lifecycle.hilt.test;
+
+        import androidx.annotation.NonNull;
+        import androidx.lifecycle.SavedStateHandle;
+        import androidx.lifecycle.hilt.ViewModelAssistedFactory;
+        import java.lang.Override;
+        import $GENERATED_TYPE;
+        import javax.inject.Inject;
+
+        $GENERATED_ANNOTATION
+        public final class MyViewModel_AssistedFactory implements
+                ViewModelAssistedFactory<MyViewModel> {
+
+            @Inject
+            MyViewModel_AssistedFactory() { }
+
+            @Override
+            @NonNull
+            public MyViewModel create(@NonNull SavedStateHandle handle) {
+                return new MyViewModel();
+            }
+        }
+        """.toJFO("androidx.lifecycle.hilt.test.MyViewModel_AssistedFactory")
+
+        val compilation = compiler()
+            .compile(myViewModel, Sources.VIEW_MODEL, Sources.SAVED_STATE_HANDLE)
+        assertThat(compilation).succeeded()
+        assertThat(compilation)
+            .generatedSourceFile("androidx.lifecycle.hilt.test.MyViewModel_AssistedFactory")
+            .hasSourceEquivalentTo(expected)
+    }
+
+    @Test
+    fun verifyAssistedFactory_savedStateOnlyArg() {
+        val myViewModel = """
+        package androidx.lifecycle.hilt.test;
+
+        import androidx.lifecycle.ViewModel;
+        import androidx.lifecycle.SavedStateHandle;
+        import androidx.lifecycle.hilt.ViewModelInject;
+
+        class MyViewModel extends ViewModel {
+            @ViewModelInject
+            MyViewModel(SavedStateHandle savedState) { }
+        }
+        """.toJFO("androidx.lifecycle.hilt.test.MyViewModel")
+
+        val expected = """
+        package androidx.lifecycle.hilt.test;
+
+        import androidx.annotation.NonNull;
+        import androidx.lifecycle.SavedStateHandle;
+        import androidx.lifecycle.hilt.ViewModelAssistedFactory;
+        import java.lang.Override;
+        import $GENERATED_TYPE;
+        import javax.inject.Inject;
+
+        $GENERATED_ANNOTATION
+        public final class MyViewModel_AssistedFactory implements
+                ViewModelAssistedFactory<MyViewModel> {
+
+            @Inject
+            MyViewModel_AssistedFactory() { }
+
+            @Override
+            @NonNull
+            public MyViewModel create(@NonNull SavedStateHandle handle) {
+                return new MyViewModel(handle);
+            }
+        }
+        """.toJFO("androidx.lifecycle.hilt.test.MyViewModel_AssistedFactory")
+
+        val compilation = compiler()
+            .compile(myViewModel, Sources.VIEW_MODEL, Sources.SAVED_STATE_HANDLE)
+        assertThat(compilation).succeeded()
+        assertThat(compilation)
+            .generatedSourceFile("androidx.lifecycle.hilt.test.MyViewModel_AssistedFactory")
+            .hasSourceEquivalentTo(expected)
+    }
+
+    @Test
+    fun verifyAssistedFactory_mixedArgs() {
+        val foo = """
+        package androidx.lifecycle.hilt.test;
+
+        public class Foo { }
+        """.toJFO("androidx.lifecycle.hilt.test.Foo")
+
+        val myViewModel = """
+        package androidx.lifecycle.hilt.test;
+
+        import androidx.lifecycle.ViewModel;
+        import androidx.lifecycle.SavedStateHandle;
+        import androidx.lifecycle.hilt.ViewModelInject;
+        import java.lang.String;
+
+        class MyViewModel extends ViewModel {
+            @ViewModelInject
+            MyViewModel(String s, Foo f, SavedStateHandle savedState, long l) { }
+        }
+        """.toJFO("androidx.lifecycle.hilt.test.MyViewModel")
+
+        val expected = """
+        package androidx.lifecycle.hilt.test;
+
+        import androidx.annotation.NonNull;
+        import androidx.lifecycle.SavedStateHandle;
+        import androidx.lifecycle.hilt.ViewModelAssistedFactory;
+        import java.lang.Long;
+        import java.lang.Override;
+        import java.lang.String;
+        import $GENERATED_TYPE;
+        import javax.inject.Inject;
+        import javax.inject.Provider;
+
+        $GENERATED_ANNOTATION
+        public final class MyViewModel_AssistedFactory implements
+                ViewModelAssistedFactory<MyViewModel> {
+
+            private final Provider<String> s;
+            private final Provider<Foo> f;
+            private final Provider<Long> l;
+
+            @Inject
+            MyViewModel_AssistedFactory(Provider<String> s, Provider<Foo> f, Provider<Long> l) {
+                this.s = s;
+                this.f = f;
+                this.l = l;
+            }
+
+            @Override
+            @NonNull
+            public MyViewModel create(@NonNull SavedStateHandle handle) {
+                return new MyViewModel(s.get(), f.get(), handle, l.get());
+            }
+        }
+        """.toJFO("androidx.lifecycle.hilt.test.MyViewModel_AssistedFactory")
+
+        val compilation = compiler()
+            .compile(foo, myViewModel, Sources.VIEW_MODEL, Sources.SAVED_STATE_HANDLE)
+        assertThat(compilation).succeeded()
+        assertThat(compilation)
+            .generatedSourceFile("androidx.lifecycle.hilt.test.MyViewModel_AssistedFactory")
+            .hasSourceEquivalentTo(expected)
+    }
+
+    @Test
+    fun verifyAssistedFactory_mixedAndProviderArgs() {
+        val foo = """
+        package androidx.lifecycle.hilt.test;
+
+        public class Foo { }
+        """.toJFO("androidx.lifecycle.hilt.test.Foo")
+
+        val myViewModel = """
+        package androidx.lifecycle.hilt.test;
+
+        import androidx.lifecycle.ViewModel;
+        import androidx.lifecycle.SavedStateHandle;
+        import androidx.lifecycle.hilt.ViewModelInject;
+        import java.lang.String;
+        import javax.inject.Provider;
+
+        class MyViewModel extends ViewModel {
+            @ViewModelInject
+            MyViewModel(String s, Provider<Foo> f, SavedStateHandle savedState) { }
+        }
+        """.toJFO("androidx.lifecycle.hilt.test.MyViewModel")
+
+        val expected = """
+        package androidx.lifecycle.hilt.test;
+
+        import androidx.annotation.NonNull;
+        import androidx.lifecycle.SavedStateHandle;
+        import androidx.lifecycle.hilt.ViewModelAssistedFactory;
+        import java.lang.Override;
+        import java.lang.String;
+        import $GENERATED_TYPE;
+        import javax.inject.Inject;
+        import javax.inject.Provider;
+
+        $GENERATED_ANNOTATION
+        public final class MyViewModel_AssistedFactory implements
+                ViewModelAssistedFactory<MyViewModel> {
+
+            private final Provider<String> s;
+            private final Provider<Foo> f;
+
+            @Inject
+            MyViewModel_AssistedFactory(Provider<String> s, Provider<Foo> f) {
+                this.s = s;
+                this.f = f;
+            }
+
+            @Override
+            @NonNull
+            public MyViewModel create(@NonNull SavedStateHandle handle) {
+                return new MyViewModel(s.get(), f, handle);
+            }
+        }
+        """.toJFO("androidx.lifecycle.hilt.test.MyViewModel_AssistedFactory")
+
+        val compilation = compiler()
+            .compile(foo, myViewModel, Sources.VIEW_MODEL, Sources.SAVED_STATE_HANDLE)
+        assertThat(compilation).succeeded()
+        assertThat(compilation)
+            .generatedSourceFile("androidx.lifecycle.hilt.test.MyViewModel_AssistedFactory")
+            .hasSourceEquivalentTo(expected)
+    }
+
+    @Test
+    fun verifyAssistedFactory_qualifiedArgs() {
+        val myQualifier = """
+        package androidx.lifecycle.hilt.test;
+
+        import javax.inject.Qualifier;
+
+        @Qualifier
+        public @interface MyQualifier { }
+        """.toJFO("androidx.lifecycle.hilt.test.MyQualifier")
+
+        val myViewModel = """
+        package androidx.lifecycle.hilt.test;
+
+        import androidx.lifecycle.ViewModel;
+        import androidx.lifecycle.SavedStateHandle;
+        import androidx.lifecycle.hilt.ViewModelInject;
+        import java.lang.Long;
+        import java.lang.String;
+        import javax.inject.Named;
+        import javax.inject.Provider;
+
+        class MyViewModel extends ViewModel {
+            @ViewModelInject
+            MyViewModel(@Named("TheString") String s, @MyQualifier Provider<Long> l,
+                    SavedStateHandle savedState) {
+            }
+        }
+        """.toJFO("androidx.lifecycle.hilt.test.MyViewModel")
+
+        val expected = """
+        package androidx.lifecycle.hilt.test;
+
+        import androidx.annotation.NonNull;
+        import androidx.lifecycle.SavedStateHandle;
+        import androidx.lifecycle.hilt.ViewModelAssistedFactory;
+        import java.lang.Long;
+        import java.lang.Override;
+        import java.lang.String;
+        import $GENERATED_TYPE;
+        import javax.inject.Inject;
+        import javax.inject.Named;
+        import javax.inject.Provider;
+
+        $GENERATED_ANNOTATION
+        public final class MyViewModel_AssistedFactory implements
+                ViewModelAssistedFactory<MyViewModel> {
+
+            private final Provider<String> s;
+            private final Provider<Long> l;
+
+            @Inject
+            MyViewModel_AssistedFactory(@Named("TheString") Provider<String> s,
+                    @MyQualifier Provider<Long> l) {
+                this.s = s;
+                this.l = l;
+            }
+
+            @Override
+            @NonNull
+            public MyViewModel create(@NonNull SavedStateHandle handle) {
+                return new MyViewModel(s.get(), l, handle);
+            }
+        }
+        """.toJFO("androidx.lifecycle.hilt.test.MyViewModel_AssistedFactory")
+
+        val compilation = compiler()
+            .compile(myQualifier, myViewModel, Sources.VIEW_MODEL, Sources.SAVED_STATE_HANDLE)
+        assertThat(compilation).succeeded()
+        assertThat(compilation)
+            .generatedSourceFile("androidx.lifecycle.hilt.test.MyViewModel_AssistedFactory")
+            .hasSourceEquivalentTo(expected)
+    }
+
+    @Test
+    fun verifyAssistedFactory_multipleSavedStateArg() {
+        val myViewModel = """
+        package androidx.lifecycle.hilt.test;
+
+        import androidx.lifecycle.ViewModel;
+        import androidx.lifecycle.SavedStateHandle;
+        import androidx.lifecycle.hilt.ViewModelInject;
+        import java.lang.String;
+
+        class MyViewModel extends ViewModel {
+            @ViewModelInject
+            MyViewModel(SavedStateHandle savedState, String s, SavedStateHandle savedState2) { }
+        }
+        """.toJFO("androidx.lifecycle.hilt.test.MyViewModel")
+
+        val expected = """
+        package androidx.lifecycle.hilt.test;
+
+        import androidx.annotation.NonNull;
+        import androidx.lifecycle.SavedStateHandle;
+        import androidx.lifecycle.hilt.ViewModelAssistedFactory;
+        import java.lang.Override;
+        import java.lang.String;
+        import $GENERATED_TYPE;
+        import javax.inject.Inject;
+        import javax.inject.Provider;
+
+        $GENERATED_ANNOTATION
+        public final class MyViewModel_AssistedFactory implements
+                ViewModelAssistedFactory<MyViewModel> {
+
+            private final Provider<String> s;
+
+            @Inject
+            MyViewModel_AssistedFactory(Provider<String> s) {
+                this.s = s;
+            }
+
+            @Override
+            @NonNull
+            public MyViewModel create(@NonNull SavedStateHandle handle) {
+                return new MyViewModel(handle, s.get(), handle);
+            }
+        }
+        """.toJFO("androidx.lifecycle.hilt.test.MyViewModel_AssistedFactory")
+
+        val compilation = compiler()
+            .compile(myViewModel, Sources.VIEW_MODEL, Sources.SAVED_STATE_HANDLE)
+        assertThat(compilation).succeeded()
+        assertThat(compilation)
+            .generatedSourceFile("androidx.lifecycle.hilt.test.MyViewModel_AssistedFactory")
+            .hasSourceEquivalentTo(expected)
+    }
+
+    @Test
+    fun verifyMultibindModule() {
+        val myViewModel = """
+        package androidx.lifecycle.hilt.test;
+
+        import androidx.lifecycle.ViewModel;
+        import androidx.lifecycle.hilt.ViewModelInject;
+
+        class MyViewModel extends ViewModel {
+            @ViewModelInject
+            MyViewModel() { }
+        }
+        """.toJFO("androidx.lifecycle.hilt.test.MyViewModel")
+
+        val expected = """
+        package androidx.lifecycle.hilt.test;
+
+        import androidx.lifecycle.ViewModel;
+        import androidx.lifecycle.hilt.ViewModelAssistedFactory;
+        import androidx.lifecycle.hilt.ViewModelKey;
+        import dagger.Binds;
+        import dagger.Module;
+        import dagger.hilt.InstallIn;
+        import dagger.hilt.android.components.ActivityRetainedComponent;
+        import dagger.multibindings.IntoMap;
+        import $GENERATED_TYPE;
+
+        $GENERATED_ANNOTATION
+        @Module
+        @InstallIn(ActivityRetainedComponent.class)
+        public interface MyViewModel_HiltModule {
+            @Binds
+            @IntoMap
+            @ViewModelKey(MyViewModel.class)
+            ViewModelAssistedFactory<? extends ViewModel> bind(MyViewModel_AssistedFactory factory)
+        }
+        """.toJFO("androidx.lifecycle.hilt.test.MyViewModel_HiltModule")
+
+        val compilation = compiler()
+            .compile(myViewModel, Sources.VIEW_MODEL, Sources.SAVED_STATE_HANDLE)
+        assertThat(compilation).succeeded()
+        assertThat(compilation)
+            .generatedSourceFile("androidx.lifecycle.hilt.test.MyViewModel_HiltModule")
+            .hasSourceEquivalentTo(expected)
+    }
+
+    @Test
+    fun verifyInnerClass() {
+        val viewModel = """
+        package androidx.lifecycle.hilt.test;
+
+        import androidx.lifecycle.ViewModel;
+        import androidx.lifecycle.hilt.ViewModelInject;
+
+        class Outer {
+            static class InnerViewModel extends ViewModel {
+                @ViewModelInject
+                InnerViewModel() { }
+            }
+        }
+        """.toJFO("androidx.lifecycle.hilt.test.Outer")
+
+        val expectedFactory = """
+        package androidx.lifecycle.hilt.test;
+
+        import androidx.annotation.NonNull;
+        import androidx.lifecycle.SavedStateHandle;
+        import androidx.lifecycle.hilt.ViewModelAssistedFactory;
+        import java.lang.Override;
+        import $GENERATED_TYPE;
+        import javax.inject.Inject;
+
+        $GENERATED_ANNOTATION
+        public final class Outer_InnerViewModel_AssistedFactory implements
+                ViewModelAssistedFactory<Outer.InnerViewModel> {
+
+            @Inject
+            Outer_InnerViewModel_AssistedFactory() { }
+
+            @Override
+            @NonNull
+            public Outer.InnerViewModel create(@NonNull SavedStateHandle handle) {
+                return new Outer.InnerViewModel();
+            }
+        }
+        """.toJFO("androidx.lifecycle.hilt.test.Outer_InnerViewModel_AssistedFactory")
+
+        val expectedModule = """
+        package androidx.lifecycle.hilt.test;
+
+        import androidx.lifecycle.ViewModel;
+        import androidx.lifecycle.hilt.ViewModelAssistedFactory;
+        import androidx.lifecycle.hilt.ViewModelKey;
+        import dagger.Binds;
+        import dagger.Module;
+        import dagger.hilt.InstallIn;
+        import dagger.hilt.android.components.ActivityRetainedComponent;
+        import dagger.multibindings.IntoMap;
+        import $GENERATED_TYPE;
+
+        $GENERATED_ANNOTATION
+        @Module
+        @InstallIn(ActivityRetainedComponent.class)
+        public interface Outer_InnerViewModel_HiltModule {
+            @Binds
+            @IntoMap
+            @ViewModelKey(Outer.InnerViewModel.class)
+            ViewModelAssistedFactory<? extends ViewModel> bind(
+                    Outer_InnerViewModel_AssistedFactory factory)
+        }
+        """.toJFO("androidx.lifecycle.hilt.test.Outer_InnerViewModel_HiltModule")
+
+        val compilation = compiler()
+            .compile(viewModel, Sources.VIEW_MODEL, Sources.SAVED_STATE_HANDLE)
+        assertThat(compilation).succeeded()
+        assertThat(compilation)
+            .generatedSourceFile("androidx.lifecycle.hilt.test" +
+                    ".Outer_InnerViewModel_AssistedFactory")
+            .hasSourceEquivalentTo(expectedFactory)
+        assertThat(compilation)
+            .generatedSourceFile("androidx.lifecycle.hilt.test" +
+                    ".Outer_InnerViewModel_HiltModule")
+            .hasSourceEquivalentTo(expectedModule)
+    }
+}
\ No newline at end of file
diff --git a/lifecycle/lifecycle-viewmodel-hilt-compiler/src/test/kotlin/androidx/lifecycle/hilt/HiltViewModelProcessorTest.kt b/lifecycle/lifecycle-viewmodel-hilt-compiler/src/test/kotlin/androidx/lifecycle/hilt/HiltViewModelProcessorTest.kt
new file mode 100644
index 00000000000..da375aa927d
--- /dev/null
+++ b/lifecycle/lifecycle-viewmodel-hilt-compiler/src/test/kotlin/androidx/lifecycle/hilt/HiltViewModelProcessorTest.kt
@@ -0,0 +1,136 @@
+/*
+ * Copyright 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 androidx.lifecycle.hilt
+
+import com.google.testing.compile.CompilationSubject.assertThat
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.junit.runners.JUnit4
+
+@RunWith(JUnit4::class)
+class HiltViewModelProcessorTest {
+
+    @Test
+    fun validViewModel() {
+        val myViewModel = """
+        package androidx.lifecycle.hilt.test;
+
+        import androidx.lifecycle.ViewModel;
+        import androidx.lifecycle.hilt.ViewModelInject;
+
+        class MyViewModel extends ViewModel {
+            @ViewModelInject
+            MyViewModel() { }
+        }
+        """.toJFO("androidx.lifecycle.hilt.test.MyViewModel")
+
+        val compilation = compiler()
+            .compile(myViewModel, Sources.VIEW_MODEL, Sources.SAVED_STATE_HANDLE)
+        assertThat(compilation).succeeded()
+    }
+
+    @Test
+    fun verifyEnclosingElementExtendsViewModel() {
+        val myViewModel = """
+        package androidx.lifecycle.hilt.test;
+
+        import androidx.lifecycle.hilt.ViewModelInject;
+
+        class MyViewModel {
+            @ViewModelInject
+            MyViewModel() { }
+        }
+        """.toJFO("androidx.lifecycle.hilt.test.MyViewModel")
+
+        val compilation = compiler().compile(myViewModel, Sources.VIEW_MODEL)
+        assertThat(compilation).failed()
+        assertThat(compilation).hadErrorCount(1)
+        assertThat(compilation)
+            .hadErrorContainingMatch("@ViewModelInject is only supported on types that subclass " +
+                    "androidx.lifecycle.ViewModel.")
+    }
+
+    @Test
+    fun verifySingleAnnotatedConstructor() {
+        val myViewModel = """
+        package androidx.lifecycle.hilt.test;
+
+        import androidx.lifecycle.ViewModel;
+        import androidx.lifecycle.hilt.ViewModelInject;
+
+        class MyViewModel extends ViewModel {
+            @ViewModelInject
+            MyViewModel() { }
+
+            @ViewModelInject
+            MyViewModel(String s) { }
+        }
+        """.toJFO("androidx.lifecycle.hilt.test.MyViewModel")
+
+        val compilation = compiler().compile(myViewModel, Sources.VIEW_MODEL)
+        assertThat(compilation).failed()
+        assertThat(compilation).hadErrorCount(1)
+        assertThat(compilation)
+            .hadErrorContainingMatch("Multiple @ViewModelInject annotated constructors found.")
+    }
+
+    @Test
+    fun verifyNonPrivateConstructor() {
+        val myViewModel = """
+        package androidx.lifecycle.hilt.test;
+
+        import androidx.lifecycle.ViewModel;
+        import androidx.lifecycle.hilt.ViewModelInject;
+
+        class MyViewModel extends ViewModel {
+            @ViewModelInject
+            private MyViewModel() { }
+        }
+        """.toJFO("androidx.lifecycle.hilt.test.MyViewModel")
+
+        val compilation = compiler().compile(myViewModel, Sources.VIEW_MODEL)
+        assertThat(compilation).failed()
+        assertThat(compilation).hadErrorCount(1)
+        assertThat(compilation)
+            .hadErrorContainingMatch("@ViewModelInject annotated constructors must not be " +
+                    "private.")
+    }
+
+    @Test
+    fun verifyInnerClassIsStatic() {
+        val myViewModel = """
+        package androidx.lifecycle.hilt.test;
+
+        import androidx.lifecycle.ViewModel;
+        import androidx.lifecycle.hilt.ViewModelInject;
+
+        class Outer {
+            class MyViewModel extends ViewModel {
+                @ViewModelInject
+                MyViewModel() { }
+            }
+        }
+        """.toJFO("androidx.lifecycle.hilt.test.Outer")
+
+        val compilation = compiler().compile(myViewModel, Sources.VIEW_MODEL)
+        assertThat(compilation).failed()
+        assertThat(compilation).hadErrorCount(1)
+        assertThat(compilation)
+            .hadErrorContainingMatch("@ViewModelInject may only be used on inner classes " +
+                    "if they are static.")
+    }
+}
\ No newline at end of file
diff --git a/lifecycle/lifecycle-viewmodel-hilt-compiler/src/test/kotlin/androidx/lifecycle/hilt/testUtils.kt b/lifecycle/lifecycle-viewmodel-hilt-compiler/src/test/kotlin/androidx/lifecycle/hilt/testUtils.kt
new file mode 100644
index 00000000000..91720a6d5c0
--- /dev/null
+++ b/lifecycle/lifecycle-viewmodel-hilt-compiler/src/test/kotlin/androidx/lifecycle/hilt/testUtils.kt
@@ -0,0 +1,42 @@
+/*
+ * Copyright 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 androidx.lifecycle.hilt
+
+import com.google.testing.compile.Compiler
+import com.google.testing.compile.Compiler.javac
+import com.google.testing.compile.JavaFileObjects
+import java.io.File
+import javax.tools.JavaFileObject
+
+object Sources {
+    val VIEW_MODEL by lazy {
+        loadJavaSource("ViewModel.java", ClassNames.VIEW_MODEL.toString())
+    }
+
+    val SAVED_STATE_HANDLE by lazy {
+        loadJavaSource("SavedStateHandle.java", ClassNames.SAVED_STATE_HANDLE.toString())
+    }
+}
+
+fun loadJavaSource(fileName: String, qName: String): JavaFileObject {
+    val contents = File("src/test/data/sources/$fileName").readText(Charsets.UTF_8)
+    return JavaFileObjects.forSourceString(qName, contents)
+}
+
+fun compiler(): Compiler = javac().withProcessors(HiltViewModelProcessor())
+
+fun String.toJFO(qName: String) = JavaFileObjects.forSourceString(qName, this.trimIndent())
\ No newline at end of file
diff --git a/lifecycle/lifecycle-viewmodel-hilt/build.gradle b/lifecycle/lifecycle-viewmodel-hilt/build.gradle
index 7e0adfdb1c2..82aae8ebe16 100644
--- a/lifecycle/lifecycle-viewmodel-hilt/build.gradle
+++ b/lifecycle/lifecycle-viewmodel-hilt/build.gradle
@@ -37,6 +37,18 @@ dependencies {
     annotationProcessor(HILT_ANDROID_COMPILER)
 }
 
+android.libraryVariants.all { variant ->
+    def name = variant.name
+    def suffix = name.capitalize()
+
+    // Create jar<variant> task for testImplementation in viewmodel-hilt--compiler.
+    project.tasks.create(name: "jar${suffix}", type: Jar){
+        dependsOn variant.javaCompileProvider.get()
+        from variant.javaCompileProvider.get().destinationDir
+        destinationDir new File(project.buildDir, "libJar")
+    }
+}
+
 androidx {
     name = "Android Lifecycle ViewModel Hilt Extension"
     publish = Publish.NONE