Merge changes Ic6152436,I7c340e73,I363add83,Icd465f45,Ia71d37ab into androidx-main

* changes:
  Run RawQueryMethodProcessorTest with KSP
  Run RemoveUnusedColumnsTest with KSP
  Add ability to pass options to processors
  Run InsertionMethodProcessorTest with KSP
  Run ShortcutMethod processor tests with KSP
diff --git a/room/compiler-processing-testing/build.gradle b/room/compiler-processing-testing/build.gradle
index e310d6e..41b9d08 100644
--- a/room/compiler-processing-testing/build.gradle
+++ b/room/compiler-processing-testing/build.gradle
@@ -74,7 +74,8 @@
 // in the source.
 tasks.named("compileTestKotlin", KotlinCompile.class).configure {
     it.kotlinOptions {
-        freeCompilerArgs += ["-Xopt-in=androidx.room.compiler.processing.ExperimentalProcessingApi"]
+        freeCompilerArgs += ["-Xopt-in=kotlin.RequiresOptIn",
+                             "-Xopt-in=androidx.room.compiler.processing.ExperimentalProcessingApi"]
     }
 }
 
diff --git a/room/compiler-processing-testing/src/main/java/androidx/room/compiler/processing/util/ProcessorTestExt.kt b/room/compiler-processing-testing/src/main/java/androidx/room/compiler/processing/util/ProcessorTestExt.kt
index d4115d2..e1c5fd5 100644
--- a/room/compiler-processing-testing/src/main/java/androidx/room/compiler/processing/util/ProcessorTestExt.kt
+++ b/room/compiler-processing-testing/src/main/java/androidx/room/compiler/processing/util/ProcessorTestExt.kt
@@ -28,6 +28,7 @@
 import com.google.common.truth.Truth.assertWithMessage
 import com.google.devtools.ksp.processing.SymbolProcessor
 import com.tschuchort.compiletesting.KotlinCompilation
+import com.tschuchort.compiletesting.kspArgs
 import com.tschuchort.compiletesting.symbolProcessors
 import java.io.ByteArrayOutputStream
 import java.io.File
@@ -79,12 +80,14 @@
 fun runProcessorTestWithoutKsp(
     sources: List<Source> = emptyList(),
     classpath: List<File> = emptyList(),
+    options: Map<String, String> = emptyMap(),
     handler: (XTestInvocation) -> Unit
 ) {
     runTests(
         params = TestCompilationParameters(
             sources = sources,
             classpath = classpath,
+            options = options,
             handlers = listOf(handler)
         ),
         JavacCompilationTestRunner,
@@ -111,8 +114,14 @@
 fun runProcessorTest(
     sources: List<Source> = emptyList(),
     classpath: List<File> = emptyList(),
+    options: Map<String, String> = emptyMap(),
     handler: (XTestInvocation) -> Unit
-) = runProcessorTest(sources = sources, classpath = classpath, handlers = listOf(handler))
+) = runProcessorTest(
+    sources = sources,
+    classpath = classpath,
+    options = options,
+    handlers = listOf(handler)
+)
 
 /**
  * Runs the step created by [createProcessingStep] with ksp and one of javac or kapt, depending
@@ -131,12 +140,14 @@
 fun runProcessorTest(
     sources: List<Source> = emptyList(),
     classpath: List<File> = emptyList(),
+    options: Map<String, String> = emptyMap(),
     createProcessingStep: () -> XProcessingStep,
     onCompilationResult: (CompilationResultSubject) -> Unit
 ) {
     runProcessorTest(
         sources = sources,
-        classpath = classpath
+        classpath = classpath,
+        options = options
     ) { invocation ->
         val step = createProcessingStep()
         val elements =
@@ -161,6 +172,7 @@
 fun runProcessorTest(
     sources: List<Source> = emptyList(),
     classpath: List<File> = emptyList(),
+    options: Map<String, String> = emptyMap(),
     handlers: List<(XTestInvocation) -> Unit>
 ) {
     val javaApRunner = if (sources.any { it is Source.KotlinSource }) {
@@ -172,6 +184,7 @@
         params = TestCompilationParameters(
             sources = sources,
             classpath = classpath,
+            options = options,
             handlers = handlers
         ),
         javaApRunner,
@@ -188,10 +201,12 @@
 fun runJavaProcessorTest(
     sources: List<Source>,
     classpath: List<File> = emptyList(),
+    options: Map<String, String> = emptyMap(),
     handler: (XTestInvocation) -> Unit
 ) = runJavaProcessorTest(
     sources = sources,
     classpath = classpath,
+    options = options,
     handlers = listOf(handler)
 )
 
@@ -202,12 +217,14 @@
 fun runJavaProcessorTest(
     sources: List<Source>,
     classpath: List<File> = emptyList(),
+    options: Map<String, String> = emptyMap(),
     handlers: List<(XTestInvocation) -> Unit>
 ) {
     runTests(
         params = TestCompilationParameters(
             sources = sources,
             classpath = classpath,
+            options = options,
             handlers = handlers
         ),
         JavacCompilationTestRunner
@@ -221,10 +238,12 @@
 fun runKaptTest(
     sources: List<Source>,
     classpath: List<File> = emptyList(),
+    options: Map<String, String> = emptyMap(),
     handler: (XTestInvocation) -> Unit
 ) = runKaptTest(
     sources = sources,
     classpath = classpath,
+    options = options,
     handlers = listOf(handler)
 )
 
@@ -235,12 +254,14 @@
 fun runKaptTest(
     sources: List<Source>,
     classpath: List<File> = emptyList(),
+    options: Map<String, String> = emptyMap(),
     handlers: List<(XTestInvocation) -> Unit>
 ) {
     runTests(
         params = TestCompilationParameters(
             sources = sources,
             classpath = classpath,
+            options = options,
             handlers = handlers
         ),
         KaptCompilationTestRunner
@@ -254,10 +275,12 @@
 fun runKspTest(
     sources: List<Source>,
     classpath: List<File> = emptyList(),
+    options: Map<String, String> = emptyMap(),
     handler: (XTestInvocation) -> Unit
 ) = runKspTest(
     sources = sources,
     classpath = classpath,
+    options = options,
     handlers = listOf(handler)
 )
 
@@ -268,12 +291,14 @@
 fun runKspTest(
     sources: List<Source>,
     classpath: List<File> = emptyList(),
+    options: Map<String, String> = emptyMap(),
     handlers: List<(XTestInvocation) -> Unit>
 ) {
     runTests(
         params = TestCompilationParameters(
             sources = sources,
             classpath = classpath,
+            options = options,
             handlers = handlers
         ),
         KspCompilationTestRunner
@@ -284,11 +309,13 @@
  * Compiles the given set of sources into a temporary folder and returns the output classes
  * directory.
  * @param sources The list of source files to compile
+ * @param options The annotation processor arguments
  * @param annotationProcessors The list of Java annotation processors to run with compilation
  * @param symbolProcessors The list of Kotlin symbol processors to run with compilation
  */
 fun compileFiles(
     sources: List<Source>,
+    options: Map<String, String> = emptyMap(),
     annotationProcessors: List<Processor> = emptyList(),
     symbolProcessors: List<SymbolProcessor> = emptyList()
 ): File {
@@ -297,6 +324,12 @@
         sources = sources,
         outputStream = outputStream
     )
+    if (annotationProcessors.isNotEmpty()) {
+        compilation.kaptArgs.putAll(options)
+    }
+    if (symbolProcessors.isNotEmpty()) {
+        compilation.kspArgs.putAll(options)
+    }
     compilation.annotationProcessors = annotationProcessors
     compilation.symbolProcessors = symbolProcessors
     val result = compilation.compile()
diff --git a/room/compiler-processing-testing/src/main/java/androidx/room/compiler/processing/util/runner/CompilationTestRunner.kt b/room/compiler-processing-testing/src/main/java/androidx/room/compiler/processing/util/runner/CompilationTestRunner.kt
index ba50be6..844bb20 100644
--- a/room/compiler-processing-testing/src/main/java/androidx/room/compiler/processing/util/runner/CompilationTestRunner.kt
+++ b/room/compiler-processing-testing/src/main/java/androidx/room/compiler/processing/util/runner/CompilationTestRunner.kt
@@ -39,5 +39,6 @@
 internal data class TestCompilationParameters(
     val sources: List<Source> = emptyList(),
     val classpath: List<File> = emptyList(),
+    val options: Map<String, String> = emptyMap(),
     val handlers: List<(XTestInvocation) -> Unit>
 )
diff --git a/room/compiler-processing-testing/src/main/java/androidx/room/compiler/processing/util/runner/JavacCompilationTestRunner.kt b/room/compiler-processing-testing/src/main/java/androidx/room/compiler/processing/util/runner/JavacCompilationTestRunner.kt
index 9097b02..d3bd3a8 100644
--- a/room/compiler-processing-testing/src/main/java/androidx/room/compiler/processing/util/runner/JavacCompilationTestRunner.kt
+++ b/room/compiler-processing-testing/src/main/java/androidx/room/compiler/processing/util/runner/JavacCompilationTestRunner.kt
@@ -48,10 +48,14 @@
         } else {
             params.sources
         }
+
+        val optionsArg = params.options.entries.map {
+            "-A${it.key}=${it.value}"
+        }
         val compiler = Compiler
             .javac()
             .withProcessors(syntheticJavacProcessor)
-            .withOptions("-Xlint")
+            .withOptions(optionsArg + "-Xlint")
             .let {
                 if (params.classpath.isNotEmpty()) {
                     it.withClasspath(params.classpath)
diff --git a/room/compiler-processing-testing/src/main/java/androidx/room/compiler/processing/util/runner/KaptCompilationTestRunner.kt b/room/compiler-processing-testing/src/main/java/androidx/room/compiler/processing/util/runner/KaptCompilationTestRunner.kt
index 6e1f860..651d73b 100644
--- a/room/compiler-processing-testing/src/main/java/androidx/room/compiler/processing/util/runner/KaptCompilationTestRunner.kt
+++ b/room/compiler-processing-testing/src/main/java/androidx/room/compiler/processing/util/runner/KaptCompilationTestRunner.kt
@@ -41,6 +41,7 @@
             outputStream = outputStream,
             classpaths = params.classpath
         )
+        compilation.kaptArgs.putAll(params.options)
         compilation.annotationProcessors = listOf(syntheticJavacProcessor)
         val result = compilation.compile()
         return KotlinCompileTestingCompilationResult(
diff --git a/room/compiler-processing-testing/src/main/java/androidx/room/compiler/processing/util/runner/KspCompilationTestRunner.kt b/room/compiler-processing-testing/src/main/java/androidx/room/compiler/processing/util/runner/KspCompilationTestRunner.kt
index 3ebaa0b..78308ba 100644
--- a/room/compiler-processing-testing/src/main/java/androidx/room/compiler/processing/util/runner/KspCompilationTestRunner.kt
+++ b/room/compiler-processing-testing/src/main/java/androidx/room/compiler/processing/util/runner/KspCompilationTestRunner.kt
@@ -25,6 +25,7 @@
 import androidx.room.compiler.processing.util.Source
 import com.tschuchort.compiletesting.KotlinCompilation
 import com.tschuchort.compiletesting.SourceFile
+import com.tschuchort.compiletesting.kspArgs
 import com.tschuchort.compiletesting.kspSourcesDir
 import com.tschuchort.compiletesting.symbolProcessors
 import java.io.ByteArrayOutputStream
@@ -57,6 +58,7 @@
             outputStream = combinedOutputStream,
             classpaths = params.classpath
         )
+        kspCompilation.kspArgs.putAll(params.options)
         kspCompilation.symbolProcessors = listOf(syntheticKspProcessor)
         kspCompilation.compile()
         // ignore KSP result for now because KSP stops compilation, which might create false
diff --git a/room/compiler-processing-testing/src/test/java/androidx/room/compiler/processing/util/TestRunnerTest.kt b/room/compiler-processing-testing/src/test/java/androidx/room/compiler/processing/util/TestRunnerTest.kt
index 7dcc46f..6ba44be 100644
--- a/room/compiler-processing-testing/src/test/java/androidx/room/compiler/processing/util/TestRunnerTest.kt
+++ b/room/compiler-processing-testing/src/test/java/androidx/room/compiler/processing/util/TestRunnerTest.kt
@@ -16,6 +16,7 @@
 
 package androidx.room.compiler.processing.util
 
+import androidx.room.compiler.processing.ExperimentalProcessingApi
 import com.google.common.truth.Truth.assertThat
 import com.squareup.javapoet.CodeBlock
 import com.squareup.javapoet.JavaFile
@@ -23,6 +24,7 @@
 import org.junit.Test
 import javax.tools.Diagnostic
 
+@OptIn(ExperimentalProcessingApi::class)
 class TestRunnerTest {
     @Test
     fun generatedBadCode_expected() = generatedBadCode(assertFailure = true)
@@ -30,6 +32,19 @@
     @Test(expected = AssertionError::class)
     fun generatedBadCode_unexpected() = generatedBadCode(assertFailure = false)
 
+    @Test
+    fun options() {
+        val testOptions = mapOf(
+            "a" to "b",
+            "c" to "d"
+        )
+        runProcessorTest(
+            options = testOptions
+        ) {
+            assertThat(it.processingEnv.options).containsAtLeastEntriesIn(testOptions)
+        }
+    }
+
     private fun generatedBadCode(assertFailure: Boolean) {
         runProcessorTest {
             if (it.processingEnv.findTypeElement("foo.Foo") == null) {
diff --git a/room/compiler/src/main/kotlin/androidx/room/processor/ProcessorErrors.kt b/room/compiler/src/main/kotlin/androidx/room/processor/ProcessorErrors.kt
index d22d0d2..386c3ec 100644
--- a/room/compiler/src/main/kotlin/androidx/room/processor/ProcessorErrors.kt
+++ b/room/compiler/src/main/kotlin/androidx/room/processor/ProcessorErrors.kt
@@ -721,6 +721,11 @@
     ) = "The partial entity $partialEntityName is missing the primary key fields " +
         "(${primaryKeyNames.joinToString()}) needed to perform an UPDATE."
 
+    fun noColumnsInPartialEntity(
+        partialEntityName: String
+    ) = "The partial entity $partialEntityName does not have any columns that can be used to " +
+        "perform the query."
+
     fun cannotFindPreparedQueryResultAdapter(
         returnType: String,
         type: QueryType
diff --git a/room/compiler/src/main/kotlin/androidx/room/processor/ShortcutMethodProcessor.kt b/room/compiler/src/main/kotlin/androidx/room/processor/ShortcutMethodProcessor.kt
index b6bddc9..59adf60 100644
--- a/room/compiler/src/main/kotlin/androidx/room/processor/ShortcutMethodProcessor.kt
+++ b/room/compiler/src/main/kotlin/androidx/room/processor/ShortcutMethodProcessor.kt
@@ -143,6 +143,15 @@
                                 ProcessorErrors.INVALID_RELATION_IN_PARTIAL_ENTITY
                             )
                         }
+
+                        if (pojo.fields.isEmpty()) {
+                            context.logger.e(
+                                executableElement,
+                                ProcessorErrors.noColumnsInPartialEntity(
+                                    partialEntityName = pojo.typeName.toString()
+                                )
+                            )
+                        }
                         onValidatePartialEntity(targetEntity, pojo)
                     }
                 }
diff --git a/room/compiler/src/test/kotlin/androidx/room/processor/InsertionMethodProcessorTest.kt b/room/compiler/src/test/kotlin/androidx/room/processor/InsertionMethodProcessorTest.kt
index f753d89..e22402f 100644
--- a/room/compiler/src/test/kotlin/androidx/room/processor/InsertionMethodProcessorTest.kt
+++ b/room/compiler/src/test/kotlin/androidx/room/processor/InsertionMethodProcessorTest.kt
@@ -21,17 +21,15 @@
 import androidx.room.Insert
 import androidx.room.OnConflictStrategy
 import androidx.room.compiler.processing.XTypeElement
+import androidx.room.compiler.processing.util.Source
+import androidx.room.compiler.processing.util.XTestInvocation
+import androidx.room.compiler.processing.util.runProcessorTest
 import androidx.room.ext.CommonTypeNames
 import androidx.room.ext.RxJava2TypeNames
 import androidx.room.ext.RxJava3TypeNames
 import androidx.room.solver.shortcut.result.InsertMethodAdapter
-import androidx.room.testing.TestInvocation
-import androidx.room.testing.TestProcessor
+import androidx.room.testing.context
 import androidx.room.vo.InsertionMethod
-import com.google.common.truth.Truth.assertAbout
-import com.google.testing.compile.CompileTester
-import com.google.testing.compile.JavaFileObjects
-import com.google.testing.compile.JavaSourcesSubjectFactory
 import com.squareup.javapoet.ArrayTypeName
 import com.squareup.javapoet.ClassName
 import com.squareup.javapoet.ParameterizedTypeName
@@ -43,8 +41,7 @@
 import org.junit.Test
 import org.junit.runner.RunWith
 import org.junit.runners.JUnit4
-import toJFO
-import javax.tools.JavaFileObject
+import toSources
 
 @Suppress("PLATFORM_CLASS_MAPPED_TO_KOTLIN")
 @RunWith(JUnit4::class)
@@ -70,14 +67,17 @@
                 @Insert
                 abstract public void foo();
                 """
-        ) { insertion, _ ->
+        ) { insertion, invocation ->
             assertThat(insertion.name, `is`("foo"))
             assertThat(insertion.parameters.size, `is`(0))
             assertThat(insertion.returnType.typeName, `is`(TypeName.VOID))
             assertThat(insertion.entities.size, `is`(0))
-        }.failsToCompile().withErrorContaining(
-            ProcessorErrors.INSERTION_DOES_NOT_HAVE_ANY_PARAMETERS_TO_INSERT
-        )
+            invocation.assertCompilationResult {
+                hasErrorContaining(
+                    ProcessorErrors.INSERTION_DOES_NOT_HAVE_ANY_PARAMETERS_TO_INSERT
+                )
+            }
+        }
     }
 
     @Test
@@ -99,7 +99,7 @@
                 `is`(ClassName.get("foo.bar", "User") as TypeName)
             )
             assertThat(insertion.returnType.typeName, `is`(TypeName.LONG))
-        }.compilesWithoutError()
+        }
     }
 
     @Test
@@ -109,13 +109,16 @@
                 @Insert
                 abstract public void foo(NotAnEntity notValid);
                 """
-        ) { insertion, _ ->
+        ) { insertion, invocation ->
             assertThat(insertion.name, `is`("foo"))
             assertThat(insertion.parameters.size, `is`(1))
             assertThat(insertion.entities.size, `is`(0))
-        }.failsToCompile().withErrorContaining(
-            ProcessorErrors.CANNOT_FIND_ENTITY_FOR_SHORTCUT_QUERY_PARAMETER
-        )
+            invocation.assertCompilationResult {
+                hasErrorContaining(
+                    ProcessorErrors.CANNOT_FIND_ENTITY_FOR_SHORTCUT_QUERY_PARAMETER
+                )
+            }
+        }
     }
 
     @Test
@@ -138,7 +141,7 @@
             assertThat(insertion.entities["u2"]?.pojo?.typeName, `is`(USER_TYPE_NAME))
             assertThat(insertion.parameters.map { it.name }, `is`(listOf("u1", "u2")))
             assertThat(insertion.returnType.typeName, `is`(TypeName.VOID))
-        }.compilesWithoutError()
+        }
     }
 
     @Test
@@ -173,7 +176,7 @@
                     ) as TypeName
                 )
             )
-        }.compilesWithoutError()
+        }
     }
 
     @Test
@@ -196,7 +199,7 @@
             assertThat(insertion.entities.size, `is`(1))
             assertThat(insertion.entities["users"]?.pojo?.typeName, `is`(USER_TYPE_NAME))
             assertThat(insertion.returnType.typeName, `is`(TypeName.VOID))
-        }.compilesWithoutError()
+        }
     }
 
     @Test
@@ -222,7 +225,7 @@
             assertThat(insertion.entities.size, `is`(1))
             assertThat(insertion.entities["users"]?.pojo?.typeName, `is`(USER_TYPE_NAME))
             assertThat(insertion.returnType.typeName, `is`(TypeName.VOID))
-        }.compilesWithoutError()
+        }
     }
 
     @Test
@@ -248,7 +251,7 @@
             assertThat(insertion.entities.size, `is`(1))
             assertThat(insertion.entities["users"]?.pojo?.typeName, `is`(USER_TYPE_NAME))
             assertThat(insertion.returnType.typeName, `is`(TypeName.VOID))
-        }.compilesWithoutError()
+        }
     }
 
     @Test
@@ -274,7 +277,7 @@
             assertThat(insertion.entities.size, `is`(1))
             assertThat(insertion.entities["users"]?.pojo?.typeName, `is`(USER_TYPE_NAME))
             assertThat(insertion.returnType.typeName, `is`(TypeName.VOID))
-        }.compilesWithoutError()
+        }
     }
 
     @Test
@@ -301,7 +304,7 @@
             assertThat(insertion.entities.size, `is`(1))
             assertThat(insertion.entities["users"]?.pojo?.typeName, `is`(USER_TYPE_NAME))
             assertThat(insertion.returnType.typeName, `is`(TypeName.VOID))
-        }.compilesWithoutError()
+        }
     }
 
     @Test
@@ -326,7 +329,7 @@
             assertThat(insertion.entities.size, `is`(2))
             assertThat(insertion.entities["u1"]?.pojo?.typeName, `is`(USER_TYPE_NAME))
             assertThat(insertion.entities["b1"]?.pojo?.typeName, `is`(BOOK_TYPE_NAME))
-        }.compilesWithoutError()
+        }
     }
 
     @Test
@@ -338,7 +341,7 @@
                 """
         ) { insertion, _ ->
             assertThat(insertion.onConflict, `is`(OnConflictStrategy.ABORT))
-        }.compilesWithoutError()
+        }
     }
 
     @Test
@@ -348,8 +351,13 @@
                 @Insert(onConflict = -1)
                 abstract public void foo(User user);
                 """
-        ) { _, _ ->
-        }.failsToCompile().withErrorContaining(ProcessorErrors.INVALID_ON_CONFLICT_VALUE)
+        ) { _, invocation ->
+            invocation.assertCompilationResult {
+                hasErrorContaining(
+                    ProcessorErrors.INVALID_ON_CONFLICT_VALUE
+                )
+            }
+        }
     }
 
     @Test
@@ -368,7 +376,7 @@
                 """
             ) { insertion, _ ->
                 assertThat(insertion.onConflict, `is`(pair.second))
-            }.compilesWithoutError()
+            }
         }
     }
 
@@ -388,11 +396,14 @@
                 @Insert
                 abstract public $type foo(User user);
                 """
-            ) { insertion, _ ->
+            ) { insertion, invocation ->
                 assertThat(insertion.methodBinder.adapter, `is`(nullValue()))
-            }.failsToCompile().withErrorContaining(
-                ProcessorErrors.CANNOT_FIND_INSERT_RESULT_ADAPTER
-            )
+                invocation.assertCompilationResult {
+                    hasErrorContaining(
+                        ProcessorErrors.CANNOT_FIND_INSERT_RESULT_ADAPTER
+                    )
+                }
+            }
         }
     }
 
@@ -410,11 +421,14 @@
                 @Insert
                 abstract public $type foo(User user);
                 """
-            ) { insertion, _ ->
+            ) { insertion, invocation ->
                 assertThat(insertion.methodBinder.adapter, `is`(nullValue()))
-            }.failsToCompile().withErrorContaining(
-                ProcessorErrors.CANNOT_FIND_INSERT_RESULT_ADAPTER
-            )
+                invocation.assertCompilationResult {
+                    hasErrorContaining(
+                        ProcessorErrors.CANNOT_FIND_INSERT_RESULT_ADAPTER
+                    )
+                }
+            }
         }
     }
 
@@ -431,11 +445,14 @@
                 @Insert
                 abstract public $type foo(User... user);
                 """
-            ) { insertion, _ ->
+            ) { insertion, invocation ->
                 assertThat(insertion.methodBinder.adapter, `is`(nullValue()))
-            }.failsToCompile().withErrorContaining(
-                ProcessorErrors.CANNOT_FIND_INSERT_RESULT_ADAPTER
-            )
+                invocation.assertCompilationResult {
+                    hasErrorContaining(
+                        ProcessorErrors.CANNOT_FIND_INSERT_RESULT_ADAPTER
+                    )
+                }
+            }
         }
     }
 
@@ -452,11 +469,14 @@
                 @Insert
                 abstract public $type foo(User user1, User user2);
                 """
-            ) { insertion, _ ->
+            ) { insertion, invocation ->
                 assertThat(insertion.methodBinder.adapter, `is`(nullValue()))
-            }.failsToCompile().withErrorContaining(
-                ProcessorErrors.CANNOT_FIND_INSERT_RESULT_ADAPTER
-            )
+                invocation.assertCompilationResult {
+                    hasErrorContaining(
+                        ProcessorErrors.CANNOT_FIND_INSERT_RESULT_ADAPTER
+                    )
+                }
+            }
         }
     }
 
@@ -526,13 +546,15 @@
                 """
             ) { insertion, _ ->
                 assertThat(insertion.methodBinder.adapter, `is`(notNullValue()))
-            }.compilesWithoutError()
+            }
         }
     }
 
     @Test
     fun targetEntitySingle() {
-        val usernameJfo = """
+        val usernameSource = Source.java(
+            "foo.bar.Username",
+            """
             package foo.bar;
             import androidx.room.*;
 
@@ -542,13 +564,14 @@
                 @ColumnInfo(name = "ageColumn")
                 int age;
             }
-        """.toJFO("foo.bar.Username")
+            """
+        )
         singleInsertMethod(
             """
                 @Insert(entity = User.class)
                 abstract public long foo(Username username);
-                """,
-            additionalJFOs = listOf(usernameJfo)
+            """,
+            additionalSources = listOf(usernameSource)
         ) { insertion, _ ->
             assertThat(insertion.name, `is`("foo"))
             assertThat(insertion.parameters.size, `is`(1))
@@ -559,7 +582,7 @@
             assertThat(insertion.entities["username"]?.isPartialEntity, `is`(true))
             assertThat(insertion.entities["username"]?.entityTypeName, `is`(USER_TYPE_NAME))
             assertThat(insertion.entities["username"]?.pojo?.typeName, `is`(USERNAME_TYPE_NAME))
-        }.compilesWithoutError()
+        }
     }
 
     @Test
@@ -568,14 +591,16 @@
             """
                 @Insert(entity = User.class)
                 abstract public long foo(User user);
-                """
+            """
         ) { _, _ ->
-        }.compilesWithoutError()
+        }
     }
 
     @Test
     fun targetEntityTwo() {
-        val usernameJfo = """
+        val usernameSource = Source.java(
+            "foo.bar.Username",
+            """
             package foo.bar;
             import androidx.room.*;
 
@@ -585,20 +610,23 @@
                 @ColumnInfo(name = "ageColumn")
                 int age;
             }
-        """.toJFO("foo.bar.Username")
+            """
+        )
         singleInsertMethod(
             """
                 @Insert(entity = User.class)
                 abstract public void foo(Username usernameA, Username usernameB);
-                """,
-            additionalJFOs = listOf(usernameJfo)
+            """,
+            additionalSources = listOf(usernameSource)
         ) { _, _ ->
-        }.compilesWithoutError()
+        }
     }
 
     @Test
     fun targetEntityMissingRequiredColumn() {
-        val usernameJfo = """
+        val usernameSource = Source.java(
+            "foo.bar.Username",
+            """
             package foo.bar;
             import androidx.room.*;
 
@@ -606,25 +634,31 @@
                 int uid;
                 String name;
             }
-        """.toJFO("foo.bar.Username")
+            """
+        )
         singleInsertMethod(
             """
                 @Insert(entity = User.class)
                 abstract public void foo(Username username);
-                """,
-            additionalJFOs = listOf(usernameJfo)
-        ) { _, _ ->
-        }.failsToCompile().withErrorContaining(
-            ProcessorErrors.missingRequiredColumnsInPartialEntity(
-                partialEntityName = USERNAME_TYPE_NAME.toString(),
-                missingColumnNames = listOf("ageColumn")
-            )
-        )
+            """,
+            additionalSources = listOf(usernameSource)
+        ) { _, invocation ->
+            invocation.assertCompilationResult {
+                hasErrorContaining(
+                    ProcessorErrors.missingRequiredColumnsInPartialEntity(
+                        partialEntityName = USERNAME_TYPE_NAME.toString(),
+                        missingColumnNames = listOf("ageColumn")
+                    )
+                )
+            }
+        }
     }
 
     @Test
     fun targetEntityColumnDefaultValue() {
-        val petNameJfo = """
+        val petNameSource = Source.java(
+            "foo.bar.PetName",
+            """
             package foo.bar;
             import androidx.room.*;
 
@@ -632,8 +666,11 @@
                 @ColumnInfo(name = "name")
                 String string;
             }
-        """.toJFO("foo.bar.PetName")
-        val petJfo = """
+            """
+        )
+        val petSource = Source.java(
+            "foo.bar.Pet",
+            """
             package foo.bar;
             import androidx.room.*;
 
@@ -645,20 +682,23 @@
                 @ColumnInfo(defaultValue = "0")
                 int age;
             }
-        """.toJFO("foo.bar.Pet")
+            """
+        )
         singleInsertMethod(
             """
                 @Insert(entity = Pet.class)
                 abstract public long foo(PetName petName);
-                """,
-            additionalJFOs = listOf(petNameJfo, petJfo)
+            """,
+            additionalSources = listOf(petNameSource, petSource)
         ) { _, _ ->
-        }.compilesWithoutError()
+        }
     }
 
     @Test
     fun targetEntityMissingPrimaryKey() {
-        val petNameJfo = """
+        val petNameSource = Source.java(
+            "foo.bar.PetName",
+            """
             package foo.bar;
             import androidx.room.*;
 
@@ -666,8 +706,11 @@
                 @ColumnInfo(name = "name")
                 String string;
             }
-        """.toJFO("foo.bar.PetName")
-        val petJfo = """
+            """
+        )
+        val petSource = Source.java(
+            "foo.bar.Pet",
+            """
             package foo.bar;
             import androidx.room.*;
 
@@ -677,25 +720,31 @@
                 int petId;
                 String name;
             }
-        """.toJFO("foo.bar.Pet")
+            """
+        )
         singleInsertMethod(
             """
                 @Insert(entity = Pet.class)
                 abstract public long foo(PetName petName);
-                """,
-            additionalJFOs = listOf(petNameJfo, petJfo)
-        ) { _, _ ->
-        }.failsToCompile().withErrorContaining(
-            ProcessorErrors.missingPrimaryKeysInPartialEntityForInsert(
-                partialEntityName = "foo.bar.PetName",
-                primaryKeyNames = listOf("petId")
-            )
-        )
+            """,
+            additionalSources = listOf(petNameSource, petSource)
+        ) { _, invocation ->
+            invocation.assertCompilationResult {
+                hasErrorContaining(
+                    ProcessorErrors.missingPrimaryKeysInPartialEntityForInsert(
+                        partialEntityName = "foo.bar.PetName",
+                        primaryKeyNames = listOf("petId")
+                    )
+                )
+            }
+        }
     }
 
     @Test
     fun targetEntityAutoGeneratedPrimaryKey() {
-        val petNameJfo = """
+        val petNameSource = Source.java(
+            "foo.bar.PetName",
+            """
             package foo.bar;
             import androidx.room.*;
 
@@ -703,8 +752,11 @@
                 @ColumnInfo(name = "name")
                 String string;
             }
-        """.toJFO("foo.bar.PetName")
-        val petJfo = """
+            """
+        )
+        val petSource = Source.java(
+            "foo.bar.Pet",
+            """
             package foo.bar;
             import androidx.room.*;
 
@@ -714,20 +766,23 @@
                 int petId;
                 String name;
             }
-        """.toJFO("foo.bar.Pet")
+            """
+        )
         singleInsertMethod(
             """
                 @Insert(entity = Pet.class)
                 abstract public long foo(PetName petName);
-                """,
-            additionalJFOs = listOf(petNameJfo, petJfo)
+            """,
+            additionalSources = listOf(petNameSource, petSource)
         ) { _, _ ->
-        }.compilesWithoutError()
+        }
     }
 
     @Test
     fun targetEntityExtraColumn() {
-        val usernameJfo = """
+        val usernameSource = Source.java(
+            "foo.bar.Username",
+            """
             package foo.bar;
             import androidx.room.*;
 
@@ -736,22 +791,28 @@
                 String name;
                 long extraField;
             }
-        """.toJFO("foo.bar.Username")
+            """
+        )
         singleInsertMethod(
             """
                 @Insert(entity = User.class)
                 abstract public long foo(Username username);
-                """,
-            additionalJFOs = listOf(usernameJfo)
-        ) { _, _ ->
-        }.failsToCompile().withErrorContaining(
-            ProcessorErrors.cannotFindAsEntityField("foo.bar.User")
-        )
+            """,
+            additionalSources = listOf(usernameSource)
+        ) { _, invocation ->
+            invocation.assertCompilationResult {
+                hasErrorContaining(
+                    ProcessorErrors.cannotFindAsEntityField("foo.bar.User")
+                )
+            }
+        }
     }
 
     @Test
     fun targetEntityExtraColumnIgnored() {
-        val usernameJfo = """
+        val usernameSource = Source.java(
+            "foo.bar.Username",
+            """
             package foo.bar;
             import androidx.room.*;
 
@@ -763,20 +824,23 @@
                 @Ignore
                 long extraField;
             }
-        """.toJFO("foo.bar.Username")
+            """
+        )
         singleInsertMethod(
             """
                 @Insert(entity = User.class)
                 abstract public long foo(Username username);
-                """,
-            additionalJFOs = listOf(usernameJfo)
+            """,
+            additionalSources = listOf(usernameSource)
         ) { _, _ ->
-        }.compilesWithoutError()
+        }
     }
 
     @Test
     fun targetEntityWithEmbedded() {
-        val usernameJfo = """
+        val usernameSource = Source.java(
+            "foo.bar.Username",
+            """
             package foo.bar;
             import androidx.room.*;
 
@@ -787,8 +851,11 @@
                 @ColumnInfo(name = "ageColumn")
                 int age;
             }
-        """.toJFO("foo.bar.Username")
-        val fullnameJfo = """
+            """
+        )
+        val fullnameSource = Source.java(
+            "foo.bar.Fullname",
+            """
             package foo.bar;
             import androidx.room.*;
 
@@ -797,20 +864,23 @@
                 String firstName;
                 String lastName;
             }
-        """.toJFO("foo.bar.Fullname")
+            """
+        )
         singleInsertMethod(
             """
                 @Insert(entity = User.class)
                 abstract public long foo(Username username);
-                """,
-            additionalJFOs = listOf(usernameJfo, fullnameJfo)
+            """,
+            additionalSources = listOf(usernameSource, fullnameSource)
         ) { _, _ ->
-        }.compilesWithoutError()
+        }
     }
 
     @Test
     fun targetEntityWithRelation() {
-        val userPetsJfo = """
+        val userPetsSource = Source.java(
+            "foo.bar.UserPets",
+            """
             package foo.bar;
             import androidx.room.*;
             import java.util.List;
@@ -820,8 +890,11 @@
                 @Relation(parentColumn = "uid", entityColumn = "ownerId")
                 List<Pet> pets;
             }
-        """.toJFO("foo.bar.UserPets")
-        val petJfo = """
+            """
+        )
+        val petSource = Source.java(
+            "foo.bar.Pet",
+            """
             package foo.bar;
             import androidx.room.*;
 
@@ -831,59 +904,59 @@
                 int petId;
                 int ownerId;
             }
-        """.toJFO("foo.bar.Pet")
+            """
+        )
         singleInsertMethod(
             """
                 @Insert(entity = User.class)
                 abstract public long foo(UserPets userPets);
                 """,
-            additionalJFOs = listOf(userPetsJfo, petJfo)
-        ) { _, _ ->
-        }.failsToCompile().withErrorContaining(ProcessorErrors.INVALID_RELATION_IN_PARTIAL_ENTITY)
+            additionalSources = listOf(userPetsSource, petSource)
+        ) { _, invocation ->
+            invocation.assertCompilationResult {
+                hasErrorContaining(
+                    ProcessorErrors.INVALID_RELATION_IN_PARTIAL_ENTITY
+                )
+            }
+        }
     }
 
     fun singleInsertMethod(
         vararg input: String,
-        additionalJFOs: List<JavaFileObject> = emptyList(),
-        handler: (InsertionMethod, TestInvocation) -> Unit
-    ): CompileTester {
-        return assertAbout(JavaSourcesSubjectFactory.javaSources())
-            .that(
-                listOf(
-                    JavaFileObjects.forSourceString(
-                        "foo.bar.MyClass",
-                        DAO_PREFIX + input.joinToString("\n") + DAO_SUFFIX
-                    ),
-                    COMMON.USER, COMMON.BOOK, COMMON.NOT_AN_ENTITY, COMMON.RX2_COMPLETABLE,
-                    COMMON.RX2_MAYBE, COMMON.RX2_SINGLE, COMMON.RX3_COMPLETABLE,
-                    COMMON.RX3_MAYBE, COMMON.RX3_SINGLE
-                ) + additionalJFOs
+        additionalSources: List<Source> = emptyList(),
+        handler: (InsertionMethod, XTestInvocation) -> Unit
+    ) {
+        val inputSource = Source.java(
+            "foo.bar.MyClass",
+            DAO_PREFIX + input.joinToString("\n") + DAO_SUFFIX
+        )
+        val commonSources = listOf(
+            COMMON.USER, COMMON.BOOK, COMMON.NOT_AN_ENTITY, COMMON.RX2_COMPLETABLE,
+            COMMON.RX2_MAYBE, COMMON.RX2_SINGLE, COMMON.RX3_COMPLETABLE,
+            COMMON.RX3_MAYBE, COMMON.RX3_SINGLE
+        ).toSources()
+
+        runProcessorTest(
+            sources = commonSources + additionalSources + inputSource
+        ) { invocation ->
+            val (owner, methods) = invocation.roundEnv
+                .getElementsAnnotatedWith(Dao::class.qualifiedName!!)
+                .filterIsInstance<XTypeElement>()
+                .map {
+                    Pair(
+                        it,
+                        it.getAllMethods().filter {
+                            it.hasAnnotation(Insert::class)
+                        }
+                    )
+                }.first { it.second.isNotEmpty() }
+            val processor = InsertionMethodProcessor(
+                baseContext = invocation.context,
+                containing = owner.type,
+                executableElement = methods.first()
             )
-            .processedWith(
-                TestProcessor.builder()
-                    .forAnnotations(Insert::class, Dao::class)
-                    .nextRunHandler { invocation ->
-                        val (owner, methods) = invocation.roundEnv
-                            .getElementsAnnotatedWith(Dao::class.qualifiedName!!)
-                            .filterIsInstance<XTypeElement>()
-                            .map {
-                                Pair(
-                                    it,
-                                    it.getAllMethods().filter {
-                                        it.hasAnnotation(Insert::class)
-                                    }
-                                )
-                            }.first { it.second.isNotEmpty() }
-                        val processor = InsertionMethodProcessor(
-                            baseContext = invocation.context,
-                            containing = owner.type,
-                            executableElement = methods.first()
-                        )
-                        val processed = processor.process()
-                        handler(processed, invocation)
-                        true
-                    }
-                    .build()
-            )
+            val processed = processor.process()
+            handler(processed, invocation)
+        }
     }
 }
diff --git a/room/compiler/src/test/kotlin/androidx/room/processor/RawQueryMethodProcessorTest.kt b/room/compiler/src/test/kotlin/androidx/room/processor/RawQueryMethodProcessorTest.kt
index 00b3b2b..001e186 100644
--- a/room/compiler/src/test/kotlin/androidx/room/processor/RawQueryMethodProcessorTest.kt
+++ b/room/compiler/src/test/kotlin/androidx/room/processor/RawQueryMethodProcessorTest.kt
@@ -17,32 +17,25 @@
 package androidx.room.processor
 
 import COMMON
-import androidx.room.ColumnInfo
 import androidx.room.Dao
-import androidx.room.Entity
-import androidx.room.PrimaryKey
-import androidx.room.Query
 import androidx.room.RawQuery
-import androidx.room.compiler.processing.util.runProcessorTest
 import androidx.room.compiler.processing.XTypeElement
+import androidx.room.compiler.processing.util.Source
+import androidx.room.compiler.processing.util.XTestInvocation
+import androidx.room.compiler.processing.util.runProcessorTest
 import androidx.room.ext.PagingTypeNames
 import androidx.room.ext.SupportDbTypeNames
 import androidx.room.processor.ProcessorErrors.RAW_QUERY_STRING_PARAMETER_REMOVED
-import androidx.room.testing.TestInvocation
-import androidx.room.testing.TestProcessor
 import androidx.room.testing.context
 import androidx.room.vo.RawQueryMethod
 import androidx.sqlite.db.SupportSQLiteQuery
-import com.google.common.truth.Truth
-import com.google.testing.compile.CompileTester
-import com.google.testing.compile.JavaFileObjects
-import com.google.testing.compile.JavaSourcesSubjectFactory
 import com.squareup.javapoet.ArrayTypeName
 import com.squareup.javapoet.ClassName
 import com.squareup.javapoet.TypeName
 import org.hamcrest.CoreMatchers.`is`
 import org.hamcrest.MatcherAssert.assertThat
 import org.junit.Test
+import toSources
 
 class RawQueryMethodProcessorTest {
     @Test
@@ -67,7 +60,7 @@
                 query.returnType.typeName,
                 `is`(ArrayTypeName.of(TypeName.INT) as TypeName)
             )
-        }.compilesWithoutError()
+        }
     }
 
     @Test
@@ -77,8 +70,11 @@
                 @RawQuery
                 abstract public int[] foo(String query);
                 """
-        ) { _, _ ->
-        }.failsToCompile().withErrorContaining(RAW_QUERY_STRING_PARAMETER_REMOVED)
+        ) { _, invocation ->
+            invocation.assertCompilationResult {
+                hasErrorContaining(RAW_QUERY_STRING_PARAMETER_REMOVED)
+            }
+        }
     }
 
     @Test
@@ -101,7 +97,7 @@
             )
             assertThat(query.observedTableNames.size, `is`(1))
             assertThat(query.observedTableNames, `is`(setOf("User")))
-        }.compilesWithoutError()
+        }
     }
 
     @Test
@@ -111,7 +107,7 @@
                 @RawQuery(observedEntities = {})
                 abstract public LiveData<User> foo(SupportSQLiteQuery query);
                 """
-        ) { query, _ ->
+        ) { query, invocation ->
             assertThat(query.name, `is`("foo"))
             assertThat(
                 query.runtimeQueryParam,
@@ -123,8 +119,12 @@
                 )
             )
             assertThat(query.observedTableNames, `is`(emptySet()))
-        }.failsToCompile()
-            .withErrorContaining(ProcessorErrors.OBSERVABLE_QUERY_NOTHING_TO_OBSERVE)
+            invocation.assertCompilationResult {
+                hasErrorContaining(
+                    ProcessorErrors.OBSERVABLE_QUERY_NOTHING_TO_OBSERVE
+                )
+            }
+        }
     }
 
     @Test
@@ -134,10 +134,13 @@
                 @RawQuery
                 abstract public ${PagingTypeNames.DATA_SOURCE_FACTORY}<Integer, User> getOne();
                 """
-        ) { _, _ ->
-            // do nothing
-        }.failsToCompile()
-            .withErrorContaining(ProcessorErrors.OBSERVABLE_QUERY_NOTHING_TO_OBSERVE)
+        ) { _, invocation ->
+            invocation.assertCompilationResult {
+                hasErrorContaining(
+                    ProcessorErrors.OBSERVABLE_QUERY_NOTHING_TO_OBSERVE
+                )
+            }
+        }
     }
 
     @Test
@@ -147,10 +150,13 @@
                 @RawQuery
                 abstract public ${PagingTypeNames.POSITIONAL_DATA_SOURCE}<User> getOne();
                 """
-        ) { _, _ ->
-            // do nothing
-        }.failsToCompile()
-            .withErrorContaining(ProcessorErrors.OBSERVABLE_QUERY_NOTHING_TO_OBSERVE)
+        ) { _, invocation ->
+            invocation.assertCompilationResult {
+                hasErrorContaining(
+                    ProcessorErrors.OBSERVABLE_QUERY_NOTHING_TO_OBSERVE
+                )
+            }
+        }
     }
 
     @Test
@@ -163,7 +169,7 @@
                 """
         ) { _, _ ->
             // do nothing
-        }.compilesWithoutError()
+        }
     }
 
     @Test
@@ -192,7 +198,7 @@
             )
             assertThat(query.returnType.typeName, `is`(pojo))
             assertThat(query.observedTableNames, `is`(emptySet()))
-        }.compilesWithoutError()
+        }
     }
 
     @Test
@@ -202,10 +208,13 @@
                 @RawQuery
                 abstract public void foo(SupportSQLiteQuery query);
                 """
-        ) { _, _ ->
-        }.failsToCompile().withErrorContaining(
-            ProcessorErrors.RAW_QUERY_BAD_RETURN_TYPE
-        )
+        ) { _, invocation ->
+            invocation.assertCompilationResult {
+                hasErrorContaining(
+                    ProcessorErrors.RAW_QUERY_BAD_RETURN_TYPE
+                )
+            }
+        }
     }
 
     interface RawQuerySuspendUnitDao {
@@ -237,10 +246,13 @@
                 @RawQuery
                 abstract public int[] foo();
                 """
-        ) { _, _ ->
-        }.failsToCompile().withErrorContaining(
-            ProcessorErrors.RAW_QUERY_BAD_PARAMS
-        )
+        ) { _, invocation ->
+            invocation.assertCompilationResult {
+                hasErrorContaining(
+                    ProcessorErrors.RAW_QUERY_BAD_PARAMS
+                )
+            }
+        }
     }
 
     @Test
@@ -251,10 +263,11 @@
                 abstract public int[] foo(SupportSQLiteQuery query,
                                           SupportSQLiteQuery query2);
                 """
-        ) { _, _ ->
-        }.failsToCompile().withErrorContaining(
-            ProcessorErrors.RAW_QUERY_BAD_PARAMS
-        )
+        ) { _, invocation ->
+            invocation.assertCompilationResult {
+                hasErrorContaining(ProcessorErrors.RAW_QUERY_BAD_PARAMS)
+            }
+        }
     }
 
     @Test
@@ -264,10 +277,13 @@
                 @RawQuery
                 abstract public int[] foo(SupportSQLiteQuery... query);
                 """
-        ) { _, _ ->
-        }.failsToCompile().withErrorContaining(
-            ProcessorErrors.RAW_QUERY_BAD_PARAMS
-        )
+        ) { _, invocation ->
+            invocation.assertCompilationResult {
+                hasErrorContaining(
+                    ProcessorErrors.RAW_QUERY_BAD_PARAMS
+                )
+            }
+        }
     }
 
     @Test
@@ -277,10 +293,13 @@
                 @RawQuery(observedEntities = {${COMMON.NOT_AN_ENTITY_TYPE_NAME}.class})
                 abstract public int[] foo(SupportSQLiteQuery query);
                 """
-        ) { _, _ ->
-        }.failsToCompile().withErrorContaining(
-            ProcessorErrors.rawQueryBadEntity(COMMON.NOT_AN_ENTITY_TYPE_NAME)
-        )
+        ) { _, invocation ->
+            invocation.assertCompilationResult {
+                hasErrorContaining(
+                    ProcessorErrors.rawQueryBadEntity(COMMON.NOT_AN_ENTITY_TYPE_NAME)
+                )
+            }
+        }
     }
 
     @Test
@@ -300,7 +319,7 @@
                 """
         ) { method, _ ->
             assertThat(method.observedTableNames, `is`(setOf("User")))
-        }.compilesWithoutError()
+        }
     }
 
     @Test
@@ -317,56 +336,46 @@
                 """
         ) { method, _ ->
             assertThat(method.observedTableNames, `is`(setOf("User")))
-        }.compilesWithoutError()
+        }
     }
 
     private fun singleQueryMethod(
         vararg input: String,
-        handler: (RawQueryMethod, TestInvocation) -> Unit
-    ): CompileTester {
-        return Truth.assertAbout(JavaSourcesSubjectFactory.javaSources())
-            .that(
-                listOf(
-                    JavaFileObjects.forSourceString(
-                        "foo.bar.MyClass",
-                        DAO_PREFIX +
-                            input.joinToString("\n") +
-                            DAO_SUFFIX
-                    ),
-                    COMMON.LIVE_DATA, COMMON.COMPUTABLE_LIVE_DATA, COMMON.USER,
-                    COMMON.DATA_SOURCE_FACTORY, COMMON.POSITIONAL_DATA_SOURCE,
-                    COMMON.NOT_AN_ENTITY
-                )
-            )
-            .processedWith(
-                TestProcessor.builder()
-                    .forAnnotations(
-                        Query::class, Dao::class, ColumnInfo::class,
-                        Entity::class, PrimaryKey::class, RawQuery::class
+        handler: (RawQueryMethod, XTestInvocation) -> Unit
+    ) {
+        val inputSource = Source.java(
+            "foo.bar.MyClass",
+            DAO_PREFIX +
+                input.joinToString("\n") +
+                DAO_SUFFIX
+        )
+        val commonSources = listOf(
+            COMMON.LIVE_DATA, COMMON.COMPUTABLE_LIVE_DATA, COMMON.USER,
+            COMMON.DATA_SOURCE_FACTORY, COMMON.POSITIONAL_DATA_SOURCE,
+            COMMON.NOT_AN_ENTITY
+        ).toSources()
+        runProcessorTest(
+            sources = commonSources + inputSource
+        ) { invocation ->
+            val (owner, methods) = invocation.roundEnv
+                .getElementsAnnotatedWith(Dao::class.qualifiedName!!)
+                .filterIsInstance<XTypeElement>()
+                .map {
+                    Pair(
+                        it,
+                        it.getAllMethods().filter {
+                            it.hasAnnotation(RawQuery::class)
+                        }
                     )
-                    .nextRunHandler { invocation ->
-                        val (owner, methods) = invocation.roundEnv
-                            .getElementsAnnotatedWith(Dao::class.qualifiedName!!)
-                            .filterIsInstance<XTypeElement>()
-                            .map {
-                                Pair(
-                                    it,
-                                    it.getAllMethods().filter {
-                                        it.hasAnnotation(RawQuery::class)
-                                    }
-                                )
-                            }.first { it.second.isNotEmpty() }
-                        val parser = RawQueryMethodProcessor(
-                            baseContext = invocation.context,
-                            containing = owner.type,
-                            executableElement = methods.first()
-                        )
-                        val parsedQuery = parser.process()
-                        handler(parsedQuery, invocation)
-                        true
-                    }
-                    .build()
+                }.first { it.second.isNotEmpty() }
+            val parser = RawQueryMethodProcessor(
+                baseContext = invocation.context,
+                containing = owner.type,
+                executableElement = methods.first()
             )
+            val parsedQuery = parser.process()
+            handler(parsedQuery, invocation)
+        }
     }
 
     companion object {
diff --git a/room/compiler/src/test/kotlin/androidx/room/processor/RemoveUnusedColumnsTest.kt b/room/compiler/src/test/kotlin/androidx/room/processor/RemoveUnusedColumnsTest.kt
index d030b1b..76bcdd6 100644
--- a/room/compiler/src/test/kotlin/androidx/room/processor/RemoveUnusedColumnsTest.kt
+++ b/room/compiler/src/test/kotlin/androidx/room/processor/RemoveUnusedColumnsTest.kt
@@ -17,47 +17,51 @@
 package androidx.room.processor
 
 import COMMON
+import androidx.room.DatabaseProcessingStep
 import androidx.room.RewriteQueriesToDropUnusedColumns
-import androidx.room.RoomProcessor
-import com.google.common.truth.Truth.assertAbout
-import com.google.testing.compile.CompileTester
-import com.google.testing.compile.JavaFileObjects
-import com.google.testing.compile.JavaSourcesSubjectFactory
+import androidx.room.compiler.processing.util.CompilationResultSubject
+import androidx.room.compiler.processing.util.Source
+import androidx.room.compiler.processing.util.runProcessorTest
 import org.junit.Test
 import org.junit.runner.RunWith
 import org.junit.runners.JUnit4
-import javax.tools.JavaFileObject
-import javax.tools.StandardLocation
 
 @RunWith(JUnit4::class)
 class RemoveUnusedColumnsTest {
 
     @Test
     fun noAnnotationGivesWarning() {
-        compile()
-            .withWarningCount(1)
-            .withWarningContaining("The query returns some columns [uid, ageColumn]")
+        compile { result ->
+            result.hasWarningContaining("The query returns some columns [uid, ageColumn]")
+            result.hasWarningCount(1)
+        }
     }
 
     @Test
     fun annotateMethod() {
         compile(
             annotateMethod = true
-        ).withWarningCount(0)
+        ) { result ->
+            result.hasNoWarnings()
+        }
     }
 
     @Test
     fun annotateDao() {
         compile(
             annotateDao = true
-        ).withWarningCount(0)
+        ) { result ->
+            result.hasNoWarnings()
+        }
     }
 
     @Test
     fun annotateDb() {
         compile(
             annotateDb = true
-        ).withWarningCount(0)
+        ) { result ->
+            result.hasNoWarnings()
+        }
     }
 
     @Test
@@ -65,8 +69,10 @@
         compile(
             annotateDb = true,
             enableExpandProjection = true
-        ).withWarningCount(1)
-            .withWarningContaining(ProcessorErrors.EXPAND_PROJECTION_ALONG_WITH_REMOVE_UNUSED)
+        ) { result ->
+            result.hasWarningContaining(ProcessorErrors.EXPAND_PROJECTION_ALONG_WITH_REMOVE_UNUSED)
+            result.hasWarningCount(1)
+        }
     }
 
     @Test
@@ -74,8 +80,10 @@
         compile(
             annotateMethod = true,
             enableExpandProjection = true
-        ).withWarningCount(1)
-            .withWarningContaining(ProcessorErrors.EXPAND_PROJECTION_ALONG_WITH_REMOVE_UNUSED)
+        ) { result ->
+            result.hasWarningContaining(ProcessorErrors.EXPAND_PROJECTION_ALONG_WITH_REMOVE_UNUSED)
+            result.hasWarningCount(1)
+        }
     }
 
     @Test
@@ -83,41 +91,44 @@
         compile(
             annotateDao = true,
             enableExpandProjection = true
-        ).withWarningCount(1)
-            .withWarningContaining(ProcessorErrors.EXPAND_PROJECTION_ALONG_WITH_REMOVE_UNUSED)
+        ) { result ->
+            result.hasWarningContaining(
+                ProcessorErrors.EXPAND_PROJECTION_ALONG_WITH_REMOVE_UNUSED
+            )
+            result.hasWarningCount(1)
+        }
     }
 
     private fun compile(
         annotateDb: Boolean = false,
         annotateDao: Boolean = false,
         annotateMethod: Boolean = false,
-        enableExpandProjection: Boolean = false
-    ): CompileTester.SuccessfulCompilationClause {
-        val jfos = dao(
+        enableExpandProjection: Boolean = false,
+        validate: (CompilationResultSubject) -> Unit,
+    ) {
+        val sources = dao(
             annotateDao = annotateDao,
             annotateDb = annotateDb,
             annotateMethod = annotateMethod
-        ) + COMMON.USER
-        return assertAbout(JavaSourcesSubjectFactory.javaSources())
-            .that(jfos)
-            .withCompilerOptions("-Xlint:-processing") // remove unclaimed annotation warnings
-            .also {
-                if (enableExpandProjection) {
-                    it.withCompilerOptions("-Aroom.expandProjection=true")
-                }
-            }
-            .processedWith(RoomProcessor())
-            .compilesWithoutError()
-            .also {
-                it.and()
-                    .generatesFileNamed(
-                        StandardLocation.CLASS_OUTPUT, "foo.bar", "MyDao_Impl.class"
-                    )
-                    .and()
-                    .generatesFileNamed(
-                        StandardLocation.CLASS_OUTPUT, "foo.bar", "MyDb_Impl.class"
-                    )
-            }
+        ) + Source.fromJavaFileObject(COMMON.USER)
+
+        runProcessorTest(
+            sources = sources,
+            createProcessingStep = {
+                DatabaseProcessingStep()
+            },
+            options = mapOf(
+                "room.expandProjection" to enableExpandProjection.toString()
+            )
+        ) { result ->
+            validate(result)
+            result.generatedSourceFileWithPath(
+                "foo/bar/MyDao_Impl.java"
+            )
+            result.generatedSourceFileWithPath(
+                "foo/bar/MyDb_Impl.java"
+            )
+        }
     }
 
     companion object {
@@ -125,14 +136,14 @@
             annotateDb: Boolean,
             annotateDao: Boolean,
             annotateMethod: Boolean
-        ): List<JavaFileObject> {
+        ): List<Source> {
             fun annotationText(enabled: Boolean) = if (enabled) {
                 "@${RewriteQueriesToDropUnusedColumns::class.java.canonicalName}"
             } else {
                 ""
             }
 
-            val pojo = JavaFileObjects.forSourceString(
+            val pojo = Source.java(
                 "foo.bar.Pojo",
                 """
                     package foo.bar;
@@ -142,7 +153,7 @@
                     }
                 """.trimIndent()
             )
-            val dao = JavaFileObjects.forSourceString(
+            val dao = Source.java(
                 "foo.bar.MyDao",
                 """
                     package foo.bar;
@@ -156,7 +167,7 @@
                     }
                 """.trimIndent()
             )
-            val db = JavaFileObjects.forSourceString(
+            val db = Source.java(
                 "foo.bar.MyDb",
                 """
                     package foo.bar;
diff --git a/room/compiler/src/test/kotlin/androidx/room/processor/ShortcutMethodProcessorTest.kt b/room/compiler/src/test/kotlin/androidx/room/processor/ShortcutMethodProcessorTest.kt
index 4621ccb..a105248 100644
--- a/room/compiler/src/test/kotlin/androidx/room/processor/ShortcutMethodProcessorTest.kt
+++ b/room/compiler/src/test/kotlin/androidx/room/processor/ShortcutMethodProcessorTest.kt
@@ -18,20 +18,18 @@
 
 import COMMON
 import androidx.room.Dao
+import androidx.room.compiler.processing.XMethodElement
+import androidx.room.compiler.processing.XType
+import androidx.room.compiler.processing.XTypeElement
+import androidx.room.compiler.processing.util.Source
+import androidx.room.compiler.processing.util.XTestInvocation
+import androidx.room.compiler.processing.util.runProcessorTest
 import androidx.room.ext.CommonTypeNames
 import androidx.room.ext.GuavaUtilConcurrentTypeNames
 import androidx.room.ext.RxJava2TypeNames
 import androidx.room.ext.RxJava3TypeNames
-import androidx.room.compiler.processing.XMethodElement
-import androidx.room.compiler.processing.XType
-import androidx.room.compiler.processing.XTypeElement
-import androidx.room.testing.TestInvocation
-import androidx.room.testing.TestProcessor
+import androidx.room.testing.context
 import androidx.room.vo.ShortcutMethod
-import com.google.common.truth.Truth
-import com.google.testing.compile.CompileTester
-import com.google.testing.compile.JavaFileObjects
-import com.google.testing.compile.JavaSourcesSubjectFactory
 import com.squareup.javapoet.ArrayTypeName
 import com.squareup.javapoet.ClassName
 import com.squareup.javapoet.ParameterizedTypeName
@@ -39,8 +37,7 @@
 import org.hamcrest.CoreMatchers.`is`
 import org.hamcrest.MatcherAssert.assertThat
 import org.junit.Test
-import toJFO
-import javax.tools.JavaFileObject
+import toSources
 import kotlin.reflect.KClass
 
 /**
@@ -70,10 +67,13 @@
                 @${annotation.java.canonicalName}
                 abstract public void foo();
                 """
-        ) { shortcut, _ ->
+        ) { shortcut, invocation ->
             assertThat(shortcut.name, `is`("foo"))
             assertThat(shortcut.parameters.size, `is`(0))
-        }.failsToCompile().withErrorContaining(noParamsError())
+            invocation.assertCompilationResult {
+                hasErrorContaining(noParamsError())
+            }
+        }
     }
 
     abstract fun noParamsError(): String
@@ -94,7 +94,7 @@
             assertThat(shortcut.entities.size, `is`(1))
             assertThat(shortcut.entities["user"]?.isPartialEntity, `is`(false))
             assertThat(shortcut.entities["user"]?.pojo?.typeName, `is`(USER_TYPE_NAME))
-        }.compilesWithoutError()
+        }
     }
 
     @Test
@@ -104,13 +104,16 @@
                 @${annotation.java.canonicalName}
                 abstract public void foo(NotAnEntity notValid);
                 """
-        ) { shortcut, _ ->
+        ) { shortcut, invocation ->
             assertThat(shortcut.name, `is`("foo"))
             assertThat(shortcut.parameters.size, `is`(1))
             assertThat(shortcut.entities.size, `is`(0))
-        }.failsToCompile().withErrorContaining(
-            ProcessorErrors.CANNOT_FIND_ENTITY_FOR_SHORTCUT_QUERY_PARAMETER
-        )
+            invocation.assertCompilationResult {
+                hasErrorContaining(
+                    ProcessorErrors.CANNOT_FIND_ENTITY_FOR_SHORTCUT_QUERY_PARAMETER
+                )
+            }
+        }
     }
 
     @Test
@@ -135,7 +138,7 @@
                 shortcut.parameters.map { it.name },
                 `is`(listOf("u1", "u2"))
             )
-        }.compilesWithoutError()
+        }
     }
 
     @Test
@@ -171,7 +174,7 @@
                 assertThat(param.pojoType?.typeName, `is`(USER_TYPE_NAME))
                 assertThat(shortcut.entities.size, `is`(1))
                 assertThat(shortcut.entities["users"]?.pojo?.typeName, `is`(USER_TYPE_NAME))
-            }.compilesWithoutError()
+            }
         }
     }
 
@@ -194,7 +197,7 @@
             )
             assertThat(shortcut.entities.size, `is`(1))
             assertThat(shortcut.entities["users"]?.pojo?.typeName, `is`(USER_TYPE_NAME))
-        }.compilesWithoutError()
+        }
     }
 
     @Test
@@ -219,7 +222,7 @@
             )
             assertThat(shortcut.entities.size, `is`(1))
             assertThat(shortcut.entities["users"]?.pojo?.typeName, `is`(USER_TYPE_NAME))
-        }.compilesWithoutError()
+        }
     }
 
     @Test
@@ -244,7 +247,7 @@
             )
             assertThat(shortcut.entities.size, `is`(1))
             assertThat(shortcut.entities["users"]?.pojo?.typeName, `is`(USER_TYPE_NAME))
-        }.compilesWithoutError()
+        }
     }
 
     @Test
@@ -270,7 +273,7 @@
             )
             assertThat(shortcut.entities.size, `is`(1))
             assertThat(shortcut.entities["users"]?.pojo?.typeName, `is`(USER_TYPE_NAME))
-        }.compilesWithoutError()
+        }
     }
 
     @Test
@@ -306,7 +309,7 @@
                 assertThat(shortcut.entities.size, `is`(2))
                 assertThat(shortcut.entities["u1"]?.pojo?.typeName, `is`(USER_TYPE_NAME))
                 assertThat(shortcut.entities["b1"]?.pojo?.typeName, `is`(BOOK_TYPE_NAME))
-            }.compilesWithoutError()
+            }
         }
     }
 
@@ -331,14 +334,19 @@
                 @${annotation.java.canonicalName}
                 abstract public $type foo(User user);
                 """
-            ) { _, _ ->
-            }.failsToCompile().withErrorContaining(invalidReturnTypeError())
+            ) { _, invocation ->
+                invocation.assertCompilationResult {
+                    hasErrorContaining(invalidReturnTypeError())
+                }
+            }
         }
     }
 
     @Test
     fun targetEntity() {
-        val usernameJfo = """
+        val usernameSource = Source.java(
+            "foo.bar.Username",
+            """
             package foo.bar;
             import androidx.room.*;
 
@@ -346,13 +354,14 @@
                 int uid;
                 String name;
             }
-        """.toJFO("foo.bar.Username")
+            """
+        )
         singleShortcutMethod(
             """
                 @${annotation.java.canonicalName}(entity = User.class)
                 abstract public int foo(Username username);
                 """,
-            additionalJFOs = listOf(usernameJfo)
+            additionalSources = listOf(usernameSource)
         ) { shortcut, _ ->
             assertThat(shortcut.name, `is`("foo"))
             assertThat(shortcut.parameters.size, `is`(1))
@@ -363,7 +372,7 @@
             assertThat(shortcut.entities["username"]?.isPartialEntity, `is`(true))
             assertThat(shortcut.entities["username"]?.entityTypeName, `is`(USER_TYPE_NAME))
             assertThat(shortcut.entities["username"]?.pojo?.typeName, `is`(USERNAME_TYPE_NAME))
-        }.compilesWithoutError()
+        }
     }
 
     @Test
@@ -374,12 +383,14 @@
                 abstract public int foo(User user);
                 """
         ) { _, _ ->
-        }.compilesWithoutError()
+        }
     }
 
     @Test
     fun targetEntityExtraColumn() {
-        val usernameJfo = """
+        val usernameSource = Source.java(
+            "foo.bar.Username",
+            """
             package foo.bar;
             import androidx.room.*;
 
@@ -388,22 +399,28 @@
                 String name;
                 long extraField;
             }
-        """.toJFO("foo.bar.Username")
+            """
+        )
         singleShortcutMethod(
             """
                 @${annotation.java.canonicalName}(entity = User.class)
                 abstract public int foo(Username username);
                 """,
-            additionalJFOs = listOf(usernameJfo)
-        ) { _, _ ->
-        }.failsToCompile().withErrorContaining(
-            ProcessorErrors.cannotFindAsEntityField("foo.bar.User")
-        )
+            additionalSources = listOf(usernameSource)
+        ) { _, invocation ->
+            invocation.assertCompilationResult {
+                hasErrorContaining(
+                    ProcessorErrors.cannotFindAsEntityField("foo.bar.User")
+                )
+            }
+        }
     }
 
     @Test
     fun targetEntityExtraColumnIgnored() {
-        val usernameJfo = """
+        val usernameSource = Source.java(
+            "foo.bar.Username",
+            """
             package foo.bar;
             import androidx.room.*;
 
@@ -413,20 +430,23 @@
                 @Ignore
                 long extraField;
             }
-        """.toJFO("foo.bar.Username")
+            """
+        )
         singleShortcutMethod(
             """
                 @${annotation.java.canonicalName}(entity = User.class)
                 abstract public int foo(Username username);
                 """,
-            additionalJFOs = listOf(usernameJfo)
+            additionalSources = listOf(usernameSource)
         ) { _, _ ->
-        }.compilesWithoutError()
+        }
     }
 
     @Test
     fun targetEntityWithEmbedded() {
-        val usernameJfo = """
+        val usernameSource = Source.java(
+            "foo.bar.Username",
+            """
             package foo.bar;
             import androidx.room.*;
 
@@ -435,8 +455,11 @@
                 @Embedded
                 Fullname name;
             }
-        """.toJFO("foo.bar.Username")
-        val fullnameJfo = """
+            """
+        )
+        val fullnameSource = Source.java(
+            "foo.bar.Fullname",
+            """
             package foo.bar;
             import androidx.room.*;
 
@@ -445,20 +468,23 @@
                 String firstName;
                 String lastName;
             }
-        """.toJFO("foo.bar.Fullname")
+            """
+        )
         singleShortcutMethod(
             """
                 @${annotation.java.canonicalName}(entity = User.class)
                 abstract public int foo(Username username);
                 """,
-            additionalJFOs = listOf(usernameJfo, fullnameJfo)
+            additionalSources = listOf(usernameSource, fullnameSource)
         ) { _, _ ->
-        }.compilesWithoutError()
+        }
     }
 
     @Test
     fun targetEntityWithRelation() {
-        val userPetsJfo = """
+        val userPetsSource = Source.java(
+            "foo.bar.UserPets",
+            """
             package foo.bar;
             import androidx.room.*;
             import java.util.List;
@@ -468,8 +494,11 @@
                 @Relation(parentColumn = "uid", entityColumn = "ownerId")
                 List<Pet> pets;
             }
-        """.toJFO("foo.bar.UserPets")
-        val petJfo = """
+            """
+        )
+        val petSource = Source.java(
+            "foo.bar.Pet",
+            """
             package foo.bar;
             import androidx.room.*;
 
@@ -479,15 +508,19 @@
                 int petId;
                 int ownerId;
             }
-        """.toJFO("foo.bar.Pet")
+            """
+        )
         singleShortcutMethod(
             """
                 @${annotation.java.canonicalName}(entity = User.class)
                 abstract public int foo(UserPets userPets);
                 """,
-            additionalJFOs = listOf(userPetsJfo, petJfo)
-        ) { _, _ ->
-        }.failsToCompile().withErrorContaining(ProcessorErrors.INVALID_RELATION_IN_PARTIAL_ENTITY)
+            additionalSources = listOf(userPetsSource, petSource)
+        ) { _, invocation ->
+            invocation.assertCompilationResult {
+                hasErrorContaining(ProcessorErrors.INVALID_RELATION_IN_PARTIAL_ENTITY)
+            }
+        }
     }
 
     @Test
@@ -496,13 +529,52 @@
             """
                 @${annotation.java.canonicalName}(entity = User.class)
                 abstract public int foo(long x);
-                """
-        ) { _, _ ->
-        }.failsToCompile().withErrorContaining(
-            ProcessorErrors.shortcutMethodArgumentMustBeAClass(
-                TypeName.LONG
-            )
+                """,
+        ) { _, invocation ->
+            invocation.assertCompilationResult {
+                if (invocation.isKsp) {
+                    hasErrorContaining(
+                        ProcessorErrors.noColumnsInPartialEntity(
+                            "java.lang.Long"
+                        )
+                    )
+                } else {
+                    // javac has a different error for primitives.
+                    hasErrorContaining(
+                        ProcessorErrors.shortcutMethodArgumentMustBeAClass(
+                            TypeName.LONG
+                        )
+                    )
+                }
+            }
+        }
+    }
+
+    @Test
+    fun targetEntity_emptyClassParameter() {
+        val emptyClass = Source.java(
+            "foo.bar.EmptyClass",
+            """
+            package foo.bar;
+            public class EmptyClass {}
+            """.trimIndent()
         )
+
+        singleShortcutMethod(
+            """
+                @${annotation.java.canonicalName}(entity = User.class)
+                abstract public int foo(EmptyClass x);
+                """,
+            additionalSources = listOf(emptyClass)
+        ) { _, invocation ->
+            invocation.assertCompilationResult {
+                hasErrorContaining(
+                    ProcessorErrors.noColumnsInPartialEntity(
+                        "foo.bar.EmptyClass"
+                    )
+                )
+            }
+        }
     }
 
     abstract fun invalidReturnTypeError(): String
@@ -515,47 +587,39 @@
 
     fun singleShortcutMethod(
         vararg input: String,
-        additionalJFOs: List<JavaFileObject> = emptyList(),
-        handler: (T, TestInvocation) -> Unit
-    ):
-        CompileTester {
-            return Truth.assertAbout(JavaSourcesSubjectFactory.javaSources())
-                .that(
-                    listOf(
-                        JavaFileObjects.forSourceString(
-                            "foo.bar.MyClass",
-                            DAO_PREFIX + input.joinToString("\n") + DAO_SUFFIX
-                        ),
-                        COMMON.USER, COMMON.BOOK, COMMON.NOT_AN_ENTITY, COMMON.RX2_COMPLETABLE,
-                        COMMON.RX2_MAYBE, COMMON.RX2_SINGLE, COMMON.RX3_COMPLETABLE,
-                        COMMON.RX3_MAYBE, COMMON.RX3_SINGLE, COMMON.LISTENABLE_FUTURE,
-                        COMMON.GUAVA_ROOM
-                    ) + additionalJFOs
-                )
-                .processedWith(
-                    TestProcessor.builder()
-                        .forAnnotations(annotation, Dao::class)
-                        .nextRunHandler { invocation ->
-                            val (owner, methods) = invocation.roundEnv
-                                .getElementsAnnotatedWith(Dao::class.qualifiedName!!)
-                                .filterIsInstance<XTypeElement>()
-                                .map {
-                                    Pair(
-                                        it,
-                                        it.getAllMethods().filter {
-                                            it.hasAnnotation(annotation)
-                                        }
-                                    )
-                                }.first { it.second.isNotEmpty() }
-                            val processed = process(
-                                baseContext = invocation.context,
-                                containing = owner.type,
-                                executableElement = methods.first()
-                            )
-                            handler(processed, invocation)
-                            true
+        additionalSources: List<Source> = emptyList(),
+        handler: (T, XTestInvocation) -> Unit
+    ) {
+        val inputSource = Source.java(
+            "foo.bar.MyClass",
+            DAO_PREFIX + input.joinToString("\n") + DAO_SUFFIX
+        )
+        val commonSources = listOf(
+            COMMON.USER, COMMON.BOOK, COMMON.NOT_AN_ENTITY, COMMON.RX2_COMPLETABLE,
+            COMMON.RX2_MAYBE, COMMON.RX2_SINGLE, COMMON.RX3_COMPLETABLE,
+            COMMON.RX3_MAYBE, COMMON.RX3_SINGLE, COMMON.LISTENABLE_FUTURE,
+            COMMON.GUAVA_ROOM
+        ).toSources()
+        runProcessorTest(
+            sources = commonSources + additionalSources + inputSource
+        ) { invocation ->
+            val (owner, methods) = invocation.roundEnv
+                .getElementsAnnotatedWith(Dao::class.qualifiedName!!)
+                .filterIsInstance<XTypeElement>()
+                .map {
+                    Pair(
+                        it,
+                        it.getAllMethods().filter {
+                            it.hasAnnotation(annotation)
                         }
-                        .build()
-                )
+                    )
+                }.first { it.second.isNotEmpty() }
+            val processed = process(
+                baseContext = invocation.context,
+                containing = owner.type,
+                executableElement = methods.first()
+            )
+            handler(processed, invocation)
         }
+    }
 }
diff --git a/room/compiler/src/test/kotlin/androidx/room/processor/UpdateMethodProcessorTest.kt b/room/compiler/src/test/kotlin/androidx/room/processor/UpdateMethodProcessorTest.kt
index 88def3e..e2a264a 100644
--- a/room/compiler/src/test/kotlin/androidx/room/processor/UpdateMethodProcessorTest.kt
+++ b/room/compiler/src/test/kotlin/androidx/room/processor/UpdateMethodProcessorTest.kt
@@ -19,6 +19,7 @@
 import androidx.room.Update
 import androidx.room.compiler.processing.XMethodElement
 import androidx.room.compiler.processing.XType
+import androidx.room.compiler.processing.util.Source
 import androidx.room.processor.ProcessorErrors.CANNOT_FIND_UPDATE_RESULT_ADAPTER
 import androidx.room.processor.ProcessorErrors.UPDATE_MISSING_PARAMS
 import androidx.room.vo.UpdateMethod
@@ -27,7 +28,6 @@
 import org.junit.Test
 import org.junit.runner.RunWith
 import org.junit.runners.JUnit4
-import toJFO
 
 @Suppress("PLATFORM_CLASS_MAPPED_TO_KOTLIN")
 @RunWith(JUnit4::class)
@@ -53,7 +53,7 @@
                 """
         ) { shortcut, _ ->
             assertThat(shortcut.onConflictStrategy, `is`(OnConflictStrategy.REPLACE))
-        }.compilesWithoutError()
+        }
     }
 
     @Test
@@ -63,32 +63,41 @@
                 @Update(onConflict = -1)
                 abstract public void foo(User user);
                 """
-        ) { _, _ ->
-        }.failsToCompile().withErrorContaining(ProcessorErrors.INVALID_ON_CONFLICT_VALUE)
+        ) { _, invocation ->
+            invocation.assertCompilationResult {
+                hasErrorContaining(ProcessorErrors.INVALID_ON_CONFLICT_VALUE)
+            }
+        }
     }
 
     @Test
     fun targetEntityMissingPrimaryKey() {
-        val usernameJfo = """
+        val usernameSource = Source.java(
+            "foo.bar.Username",
+            """
             package foo.bar;
             import androidx.room.*;
 
             public class Username {
                 String name;
             }
-        """.toJFO("foo.bar.Username")
+            """
+        )
         singleShortcutMethod(
             """
                 @Update(entity = User.class)
                 abstract public int foo(Username username);
                 """,
-            additionalJFOs = listOf(usernameJfo)
-        ) { _, _ ->
-        }.failsToCompile().withErrorContaining(
-            ProcessorErrors.missingPrimaryKeysInPartialEntityForUpdate(
-                partialEntityName = "foo.bar.Username",
-                primaryKeyNames = listOf("uid")
-            )
-        )
+            additionalSources = listOf(usernameSource)
+        ) { _, invocation ->
+            invocation.assertCompilationResult {
+                hasErrorContaining(
+                    ProcessorErrors.missingPrimaryKeysInPartialEntityForUpdate(
+                        partialEntityName = "foo.bar.Username",
+                        primaryKeyNames = listOf("uid")
+                    )
+                )
+            }
+        }
     }
 }