From d36f246dcca6caa744e89171752cabba2a8d5807 Mon Sep 17 00:00:00 2001
From: keyonghan <54558023+keyonghan@users.noreply.github.com>
Date: Wed, 25 Aug 2021 20:36:02 -0700
Subject: [PATCH] Use bucket to check staging test instead of builder name
 (#88908)

---
 dev/devicelab/lib/command/upload_results.dart | 3 +++
 dev/devicelab/lib/framework/cocoon.dart       | 8 ++++----
 dev/devicelab/test/cocoon_test.dart           | 4 ++--
 3 files changed, 9 insertions(+), 6 deletions(-)

diff --git a/dev/devicelab/lib/command/upload_results.dart b/dev/devicelab/lib/command/upload_results.dart
index 19f4a3aa4b..1700981395 100644
--- a/dev/devicelab/lib/command/upload_results.dart
+++ b/dev/devicelab/lib/command/upload_results.dart
@@ -23,6 +23,7 @@ class UploadResultsCommand extends Command<void> {
     argParser.addOption('luci-builder', help: '[Flutter infrastructure] Name of the LUCI builder being run on.');
     argParser.addOption('test-status', help: 'Test status: Succeeded|Failed');
     argParser.addOption('commit-time', help: 'Commit time in UNIX timestamp');
+    argParser.addOption('builder-bucket', help: '[Flutter infrastructure] Luci builder bucket the test is running in.');
   }
 
   @override
@@ -40,6 +41,7 @@ class UploadResultsCommand extends Command<void> {
     final String? builderName = argResults!['luci-builder'] as String?;
     final String? testStatus = argResults!['test-status'] as String?;
     final String? commitTime = argResults!['commit-time'] as String?;
+    final String? builderBucket = argResults!['builder-bucket'] as String?;
 
     // Upload metrics to skia perf from test runner when `resultsPath` is specified.
     if (resultsPath != null) {
@@ -54,6 +56,7 @@ class UploadResultsCommand extends Command<void> {
       gitBranch: gitBranch,
       builderName: builderName,
       testStatus: testStatus,
+      builderBucket: builderBucket,
     );
   }
 }
diff --git a/dev/devicelab/lib/framework/cocoon.dart b/dev/devicelab/lib/framework/cocoon.dart
index 2fe5d0aab9..bd32afe9e3 100644
--- a/dev/devicelab/lib/framework/cocoon.dart
+++ b/dev/devicelab/lib/framework/cocoon.dart
@@ -90,6 +90,7 @@ class Cocoon {
     String? gitBranch,
     String? builderName,
     String? testStatus,
+    String? builderBucket,
   }) async {
     Map<String, dynamic> resultsJson = <String, dynamic>{};
     if (resultsPath != null) {
@@ -102,7 +103,7 @@ class Cocoon {
       resultsJson['NewStatus'] = testStatus;
     }
     resultsJson['TestFlaky'] = isTestFlaky ?? false;
-    if (_shouldUpdateCocoon(resultsJson)) {
+    if (_shouldUpdateCocoon(resultsJson, builderBucket: builderBucket)) {
       await retry(
         () async => _sendUpdateTaskRequest(resultsJson).timeout(Duration(seconds: requestTimeoutLimit)),
         retryIf: (Exception e) => e is SocketException || e is TimeoutException || e is ClientException,
@@ -112,10 +113,9 @@ class Cocoon {
   }
 
   /// Only post-submit tests on `master` are allowed to update in cocoon.
-  bool _shouldUpdateCocoon(Map<String, dynamic> resultJson) {
+  bool _shouldUpdateCocoon(Map<String, dynamic> resultJson, {String? builderBucket = 'prod'}) {
     const List<String> supportedBranches = <String>['master'];
-    return supportedBranches.contains(resultJson['CommitBranch']) &&
-        !resultJson['BuilderName'].toString().toLowerCase().contains('staging');
+    return supportedBranches.contains(resultJson['CommitBranch']) && builderBucket != 'staging';
   }
 
   /// Write the given parameters into an update task request and store the JSON in [resultsPath].
diff --git a/dev/devicelab/test/cocoon_test.dart b/dev/devicelab/test/cocoon_test.dart
index 097a28b185..e31fd83888 100644
--- a/dev/devicelab/test/cocoon_test.dart
+++ b/dev/devicelab/test/cocoon_test.dart
@@ -353,14 +353,14 @@ void main() {
       const String updateTaskJson = '{'
           '"CommitBranch":"master",'
           '"CommitSha":"$commitSha",'
-          '"BuilderName":"Linux_staging_test",'
+          '"BuilderName":"builderAbc",'
           '"NewStatus":"Succeeded",'
           '"ResultData":{"i":0.0,"j":0.0,"not_a_metric":"something"},'
           '"BenchmarkScoreKeys":["i","j"]}';
       fs.file(resultsPath).writeAsStringSync(updateTaskJson);
 
       // This will fail if it decided to upload results
-      await cocoon.sendTaskStatus(resultsPath: resultsPath);
+      await cocoon.sendTaskStatus(resultsPath: resultsPath, builderBucket: 'staging');
     });
   });