fix: public key bundling

packages/flutter_tools/gradle/src/main/groovy/flutter.groovy

@@ -1332,33 +1332,32 @@ class FlutterPlugin implements Plugin<Project> {
                 }
                 Task copyFlutterAssetsTask = addFlutterDeps(variant)
                 copyFlutterAssetsTask.doLast {
-                  def content = "";
+                  def yaml = new Yaml()
                   def outputDir = copyFlutterAssetsTask.destinationDir
+
+                  // Read the shorebird.yaml file into a map.
                   def shorebirdYamlFile = new File("${outputDir}/flutter_assets/shorebird.yaml")
+                  def shorebirdYamlData = yaml.load(shorebirdYamlFile.text)
+
+                  // Update the app_id to the correct flavor if one was provided.
+                  if (variant.flavorName != null && !variant.flavorName.isEmpty()) {
+                    shorebirdYamlData['app_id'] = shorebirdYamlData.flavors[variant.flavorName]
+                  }
+
+                  // Remove any flavors. This is a no-op if the flavors key doesn't exist.
+                  shorebirdYamlData.remove('flavors')
 
+                  // Add a public key if one was provided via an env var.
+                  // Ideally we'd have a way to pass the key as a parameter, but
+                  // an env var was the easiest way to get this working.
                   def shorebirdPublicKeyEnvVar = System.getenv('SHOREBIRD_PUBLIC_KEY')
                   if (shorebirdPublicKeyEnvVar != null && !shorebirdPublicKeyEnvVar.isEmpty()) {
-                    content += 'patch_public_key: ' + shorebirdPublicKeyEnvVar + '\n';
+                    shorebirdYamlData['patch_public_key'] = shorebirdPublicKeyEnvVar
                   }
 
-                  if (variant.flavorName != null && !variant.flavorName.isEmpty()) {
-                    def flavor = variant.flavorName
-                    def shorebirdYaml = new Yaml().load(shorebirdYamlFile.text)
-                    def flavorAppId = shorebirdYaml['flavors'][flavor]
-                    if (flavorAppId == null) {
-                        throw new GradleException("Cannot find app_id for ${flavor} in shorebird.yaml")
-                    }
-                    content += 'app_id: ' + flavorAppId + '\n';
-                    if (shorebirdYaml.containsKey('base_url')) {
-                        content += 'base_url: ' + shorebirdYaml['base_url'] + '\n';
-                    }
-                    if (shorebirdYaml.containsKey('auto_update')) {
-                        content += 'auto_update: ' + shorebirdYaml['auto_update'] + '\n';
-                    }
-                  }
-                  if (!content.isEmpty()) {
-                    shorebirdYamlFile.write(content)
-                  }
+                  // Write the updated map back to the shorebird.yaml file.
+                  def updatedYamlString = yaml.dumpAsMap(shorebirdYamlData)
+                  shorebirdYamlFile.write(updatedYamlString)
                 }
                 def variantOutput = variant.outputs.first()
                 def processResources = variantOutput.hasProperty(propProcessResourcesProvider) ?

packages/shorebird_tests/pubspec.yaml

@@ -8,6 +8,7 @@ environment:
 dependencies:
   archive: ^3.5.1
   path: ^1.9.0
+  yaml: ^3.1.2
 
 dev_dependencies:
   lints: ^3.0.0

packages/shorebird_tests/test/android_test.dart

@@ -11,64 +11,81 @@ void main() {
       expect(projectDirectory.shorebirdFile.existsSync(), isTrue);
     });
 
-    testWithShorebirdProject(
-      'adds the public key when passed through the environment variable',
-      (projectDirectory) async {
-        const base64PublicKey = 'public_123';
-        await projectDirectory.runFlutterBuildApk(
-          environment: {
-            'SHOREBIRD_PUBLIC_KEY': base64PublicKey,
-          },
-        );
-
-        final generatedYaml =
-            await projectDirectory.getGeneratedShorebirdYaml();
-
-        expect(generatedYaml, contains('patch_public_key: $base64PublicKey'));
-      },
-    );
-
-    group('when building with a flavor', () {
+    group('when passing the public key through the environment variable', () {
       testWithShorebirdProject(
-        'correctly changes the app id',
+        'adds the public key on top of the original file',
         (projectDirectory) async {
-          projectDirectory.addAndroidFlavors();
-          projectDirectory.addShorebirdFlavors();
+          final originalYaml = projectDirectory.shorebirdYaml;
 
-          await projectDirectory.runFlutterBuildApk(flavor: 'internal');
+          const base64PublicKey = 'public_123';
+          await projectDirectory.runFlutterBuildApk(
+            environment: {
+              'SHOREBIRD_PUBLIC_KEY': base64PublicKey,
+            },
+          );
 
           final generatedYaml =
-              await projectDirectory.getGeneratedShorebirdYaml(
-            flavor: 'internal',
+              await projectDirectory.getGeneratedShorebirdYaml();
+
+          expect(
+            generatedYaml.keys,
+            containsAll(originalYaml.keys),
           );
 
-          expect(generatedYaml, contains('app_id: internal_123'));
+          expect(
+            generatedYaml['patch_public_key'],
+            equals(base64PublicKey),
+          );
         },
       );
+    });
 
+    group('when building with a flavor', () {
       testWithShorebirdProject(
-        'correctly changes the app id and adds the public key when passed through the environment variable',
+        'correctly changes the app id',
         (projectDirectory) async {
-          const base64PublicKey = 'public_123';
           projectDirectory.addAndroidFlavors();
           projectDirectory.addShorebirdFlavors();
 
-          await projectDirectory.runFlutterBuildApk(
-            flavor: 'internal',
-            environment: {
-              'SHOREBIRD_PUBLIC_KEY': base64PublicKey,
-            },
-          );
+          await projectDirectory.runFlutterBuildApk(flavor: 'internal');
 
           final generatedYaml =
               await projectDirectory.getGeneratedShorebirdYaml(
             flavor: 'internal',
           );
 
-          expect(generatedYaml, contains('app_id: internal_123'));
-          expect(generatedYaml, contains('patch_public_key: $base64PublicKey'));
+          expect(generatedYaml['app_id'], equals('internal_123'));
         },
       );
+
+      group('when public key passed through environment variable', () {
+        testWithShorebirdProject(
+          'correctly changes the app id and adds the public key',
+          (projectDirectory) async {
+            const base64PublicKey = 'public_123';
+            projectDirectory.addAndroidFlavors();
+            projectDirectory.addShorebirdFlavors();
+
+            await projectDirectory.runFlutterBuildApk(
+              flavor: 'internal',
+              environment: {
+                'SHOREBIRD_PUBLIC_KEY': base64PublicKey,
+              },
+            );
+
+            final generatedYaml =
+                await projectDirectory.getGeneratedShorebirdYaml(
+              flavor: 'internal',
+            );
+
+            expect(generatedYaml['app_id'], equals('internal_123'));
+            expect(
+              generatedYaml['patch_public_key'],
+              equals(base64PublicKey),
+            );
+          },
+        );
+      });
     });
   });
 }

packages/shorebird_tests/test/shorebird_tests.dart

@@ -6,6 +6,7 @@ import 'package:path/path.dart' as path;
 
 import 'package:meta/meta.dart';
 import 'package:test/test.dart';
+import 'package:yaml/yaml.dart';
 
 /// This will be the path to the flutter binary housed in this flutter repository.
 ///
@@ -102,6 +103,9 @@ extension ShorebirdProjectDirectoryOnDirectory on Directory {
         path.join(this.path, 'shorebird.yaml'),
       );
 
+  YamlMap get shorebirdYaml =>
+      loadYaml(shorebirdFile.readAsStringSync()) as YamlMap;
+
   File get appGradleFile => File(
         path.join(this.path, 'android', 'app', 'build.gradle'),
       );
@@ -184,14 +188,14 @@ $flavors
         ),
       );
 
-  Future<String> getGeneratedShorebirdYaml({String? flavor}) async {
+  Future<YamlMap> getGeneratedShorebirdYaml({String? flavor}) async {
     final decodedBytes =
         ZipDecoder().decodeBytes(apkFile(flavor: flavor).readAsBytesSync());
 
     await extractArchiveToDisk(
         decodedBytes, path.join(this.path, 'apk-extracted'));
 
-    return File(
+    final yamlString = File(
       path.join(
         this.path,
         'apk-extracted',
@@ -200,5 +204,6 @@ $flavors
         'shorebird.yaml',
       ),
     ).readAsStringSync();
+    return loadYaml(yamlString) as YamlMap;
   }
 }