Merge pull request #1270 from michelle192837/resultstorefix

Fix ResultStore updater to only update relevant groups.
GoogleCloudPlatform · Mar 4, 2024 · 13cb8ad · 13cb8ad
2 parents d1617be + a572fd6
commit 13cb8ad
Show file tree

Hide file tree

Showing 2 changed files with 79 additions and 6 deletions.
diff --git a/pkg/updater/resultstore/resultstore.go b/pkg/updater/resultstore/resultstore.go
@@ -44,15 +44,23 @@ var _ updater.TargetResult = &singleActionResult{}
 // TestResultStatus represents the status of a test result.
 type TestResultStatus int64
 
+// shouldUpdate returns whether the ResultStore updater should update this test group.
+func shouldUpdate(log logrus.FieldLogger, tg *configpb.TestGroup, client *DownloadClient) bool {
+	if tg.GetResultSource().GetResultstoreConfig() == nil {
+		log.Debug("Skipping non-ResultStore group.")
+		return false
+	}
+	if client == nil {
+		log.WithField("name", tg.GetName()).Error("ResultStore update requested, but no client found.")
+		return false
+	}
+	return true
+}
+
 // Updater returns a ResultStore-based GroupUpdater, which knows how to process result data stored in ResultStore.
 func Updater(resultStoreClient *DownloadClient, gcsClient gcs.Client, groupTimeout time.Duration, write bool) updater.GroupUpdater {
 	return func(parent context.Context, log logrus.FieldLogger, client gcs.Client, tg *configpb.TestGroup, gridPath gcs.Path) (bool, error) {
-		if !tg.UseKubernetesClient && (tg.ResultSource == nil || tg.ResultSource.GetGcsConfig() == nil) {
-			log.Debug("Skipping non-kubernetes client group")
-			return false, nil
-		}
-		if resultStoreClient == nil {
-			log.WithField("name", tg.GetName()).Warn("ResultStore update requested, but no client found.")
+		if !shouldUpdate(log, tg, resultStoreClient) {
 			return false, nil
 		}
 		ctx, cancel := context.WithTimeout(parent, groupTimeout)

diff --git a/pkg/updater/resultstore/resultstore_test.go b/pkg/updater/resultstore/resultstore_test.go
@@ -3540,3 +3540,68 @@ func TestCalculateMetrics(t *testing.T) {
 		})
 	}
 }
+
+func TestShouldUpdate(t *testing.T) {
+	cases := []struct {
+		name         string
+		tg           *configpb.TestGroup
+		clientExists bool
+		want         bool
+	}{
+		{
+			name: "nil",
+			want: false,
+		},
+		{
+			name: "test group config only",
+			tg: &configpb.TestGroup{
+				ResultSource: &configpb.TestGroup_ResultSource{
+					ResultSourceConfig: &configpb.TestGroup_ResultSource_ResultstoreConfig{},
+				},
+			},
+			clientExists: false,
+			want:         false,
+		},
+		{
+			name:         "client only",
+			clientExists: true,
+			want:         false,
+		},
+		{
+			name: "client and non-ResultStore config",
+			tg: &configpb.TestGroup{
+				ResultSource: &configpb.TestGroup_ResultSource{
+					ResultSourceConfig: &configpb.TestGroup_ResultSource_GcsConfig{},
+				},
+			},
+			clientExists: false,
+			want:         false,
+		},
+		{
+			name: "basically works",
+			tg: &configpb.TestGroup{
+				ResultSource: &configpb.TestGroup_ResultSource{
+					ResultSourceConfig: &configpb.TestGroup_ResultSource_ResultstoreConfig{
+						ResultstoreConfig: &configpb.ResultStoreConfig{
+							Project: "my-gcp-project",
+						},
+					},
+				},
+			},
+			clientExists: true,
+			want:         true,
+		},
+	}
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			// Create a fake client if specified.
+			var dlClient *DownloadClient
+			if tc.clientExists {
+				dlClient = &DownloadClient{client: &fakeClient{}}
+			}
+			if got := shouldUpdate(logrus.WithField("name", tc.name), tc.tg, dlClient); tc.want != got {
+				t.Errorf("shouldUpdate(%v, clientExists = %t) got %t, want %t", tc.tg, tc.clientExists, got, tc.want)
+			}
+		})
+	}
+}