1 files changed, 3 insertions, 5 deletions
diff --git a/gradle/src/test/groovy/com/android/build/gradle/AppPluginDslTest.groovy b/gradle/src/test/groovy/com/android/build/gradle/AppPluginDslTest.groovy
index 756dfc7..bef2ec7 100644
--- a/gradle/src/test/groovy/com/android/build/gradle/AppPluginDslTest.groovy
+++ b/gradle/src/test/groovy/com/android/build/gradle/AppPluginDslTest.groovy
@@ -266,19 +266,17 @@ public class AppPluginDslTest extends BaseTest {
                                            @NonNull String testedVariantName,
                                            @NonNull Set<ApplicationVariant> variants,
                                            @NonNull Set<TestVariant> testVariants) {
-        ApplicationVariant variant = findNamedItem(variants, variantName)
-        assertNotNull(variant)
+        ApplicationVariant variant = findNamedItem(variants, variantName, "variantData")
         assertNotNull(variant.testVariant)
         assertEquals(testedVariantName, variant.testVariant.name)
-        assertEquals(variant.testVariant, findNamedItem(testVariants, testedVariantName))
+        assertEquals(variant.testVariant, findNamedItemMaybe(testVariants, testedVariantName))
         checkTasks(variant)
         checkTasks(variant.testVariant)
     }
 
     private static void checkNonTestedVariant(@NonNull String variantName,
                                               @NonNull Set<ApplicationVariant> variants) {
-        ApplicationVariant variant = findNamedItem(variants, variantName)
-        assertNotNull(variant)
+        ApplicationVariant variant = findNamedItem(variants, variantName, "variantData")
         assertNull(variant.testVariant)
         checkTasks(variant)
     }