Ensures flutter jar is added to all build types on plugin projects (#34573)

diff --git a/dev/devicelab/bin/tasks/gradle_plugin_light_apk_test.dart b/dev/devicelab/bin/tasks/gradle_plugin_light_apk_test.dart
index b2961cc..3f3b5c1 100644
--- a/dev/devicelab/bin/tasks/gradle_plugin_light_apk_test.dart
+++ b/dev/devicelab/bin/tasks/gradle_plugin_light_apk_test.dart
@@ -165,6 +165,17 @@
       });
 
       await runProjectTest((FlutterProject project) async {
+        section('gradlew assembleLocal (plugin with custom build type)');
+        await project.addCustomBuildType('local', initWith: 'debug');
+        await project.addGlobalBuildType('local', initWith: 'debug');
+        section('Add plugin');
+        await project.addPlugin('path_provider');
+        await project.getPackages();
+
+        await project.runGradleTask('assembleLocal');
+      });
+
+      await runProjectTest((FlutterProject project) async {
         section('gradlew assembleFreeDebug (product flavor)');
         await project.addProductFlavors(<String>['free']);
         await project.runGradleTask('assembleFreeDebug');
diff --git a/dev/devicelab/lib/framework/apk_utils.dart b/dev/devicelab/lib/framework/apk_utils.dart
index 3b17872..d4780f8 100644
--- a/dev/devicelab/lib/framework/apk_utils.dart
+++ b/dev/devicelab/lib/framework/apk_utils.dart
@@ -116,6 +116,42 @@
     ''');
   }
 
+  Future<void> addGlobalBuildType(String name, {String initWith}) async {
+    final File buildScript = File(
+      path.join(androidPath, 'build.gradle'),
+    );
+
+    buildScript.openWrite(mode: FileMode.append).write('''
+subprojects {
+  afterEvaluate {
+    android {
+        buildTypes {
+            $name {
+                initWith $initWith
+            }
+        }
+    }
+  }
+}
+    ''');
+  }
+
+  Future<void> addPlugin(String plugin) async {
+    final File pubspec = File(path.join(rootPath, 'pubspec.yaml'));
+    String content = await pubspec.readAsString();
+    content = content.replaceFirst(
+      '\ndependencies:\n',
+      '\ndependencies:\n  $plugin:\n',
+    );
+    await pubspec.writeAsString(content, flush: true);
+  }
+
+  Future<void> getPackages() async {
+    await inDirectory(Directory(rootPath), () async {
+      await flutter('pub', options: <String>['get']);
+    });
+  }
+
   Future<void> addProductFlavors(Iterable<String> flavors) async {
     final File buildScript = File(
       path.join(androidPath, 'app', 'build.gradle'),
diff --git a/packages/flutter_tools/gradle/flutter.gradle b/packages/flutter_tools/gradle/flutter.gradle
index dcf3d10..372a1a0 100644
--- a/packages/flutter_tools/gradle/flutter.gradle
+++ b/packages/flutter_tools/gradle/flutter.gradle
@@ -223,8 +223,16 @@
                             initWith debug
                         }
                     }
+                    
+                    pluginProject.android.buildTypes.each {
+                        def buildMode = buildModeFor(it)
+                        addFlutterJarCompileOnlyDependency(pluginProject, it.name, project.files( flutterJar ?: baseJar[buildMode] ))
+                    }
+                    pluginProject.android.buildTypes.whenObjectAdded {
+                        def buildMode = buildModeFor(it)
+                        addFlutterJarCompileOnlyDependency(pluginProject, it.name, project.files( flutterJar ?: baseJar[buildMode] ))
+                    }
                 }
-                pluginProject.afterEvaluate this.&addFlutterJarCompileOnlyDependency
             } else {
                 project.logger.error("Plugin project :$name not found. Please update settings.gradle.")
             }
@@ -292,33 +300,20 @@
         return PLATFORM_ARM32;
     }
 
-    // TODO(blasten): Clean this up.
-    private void addFlutterJarCompileOnlyDependency(Project project) {
+    private void addFlutterJarCompileOnlyDependency(Project project, String variantName, FileCollection files) {
         if (project.state.failure) {
             return
         }
-        project.dependencies {
-            if (flutterJar != null) {
-                if (project.getConfigurations().findByName("compileOnly")) {
-                    compileOnly project.files(flutterJar)
-                } else {
-                    provided project.files(flutterJar)
-                }
-            } else {
-                assert baseJar["debug"].isFile()
-                assert baseJar["profile"].isFile()
-                assert baseJar["release"].isFile()
 
-                if (project.getConfigurations().findByName("debugCompileOnly")) {
-                    debugCompileOnly project.files(baseJar["debug"])
-                    profileCompileOnly project.files(baseJar["profile"])
-                    releaseCompileOnly project.files(baseJar["release"])
-                } else {
-                    debugProvided project.files(baseJar["debug"])
-                    profileProvided project.files(baseJar["profile"])
-                    releaseProvided project.files(baseJar["release"])
-                }
+        project.dependencies {
+            String configuration;
+            if (project.getConfigurations().findByName("compileOnly")) {
+                configuration = "${variantName}CompileOnly";
+            } else {
+                configuration = "${variantName}Provided";
             }
+
+            add(configuration, files)
         }
     }