test(backend): Fix upgrade_test flakiness. Fixes #4426 (#4460)

kubeflow · Sep 3, 2020 · b1dcedc · b1dcedc
1 parent a77af2c
commit b1dcedc
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 1 deletion.
diff --git a/backend/test/integration/upgrade_test.go b/backend/test/integration/upgrade_test.go
@@ -360,6 +360,9 @@ func checkHelloWorldRunDetail(t *testing.T, runDetail *run_model.APIRunDetail) {
 	// Check runtime workflow manifest is not empty
 	assert.Contains(t, runDetail.PipelineRuntime.WorkflowManifest, "whalesay")
 
+	expectedExperimentID := test.GetExperimentIDFromAPIResourceReferences(runDetail.Run.ResourceReferences)
+	require.NotEmpty(t, expectedExperimentID)
+
 	expectedRun := &run_model.APIRun{
 		ID:          runDetail.Run.ID,
 		Name:        "hello world",
@@ -371,7 +374,7 @@ func checkHelloWorldRunDetail(t *testing.T, runDetail *run_model.APIRunDetail) {
 			WorkflowManifest: runDetail.Run.PipelineSpec.WorkflowManifest,
 		},
 		ResourceReferences: []*run_model.APIResourceReference{
-			{Key: &run_model.APIResourceKey{Type: run_model.APIResourceTypeEXPERIMENT, ID: runDetail.Run.ResourceReferences[0].Key.ID},
+			{Key: &run_model.APIResourceKey{Type: run_model.APIResourceTypeEXPERIMENT, ID: expectedExperimentID},
 				Name: "hello world experiment", Relationship: run_model.APIRelationshipOWNER,
 			},
 			{Key: &run_model.APIResourceKey{ID: runDetail.Run.PipelineSpec.PipelineID, Type: run_model.APIResourceTypePIPELINEVERSION},

diff --git a/backend/test/test_utils.go b/backend/test/test_utils.go
@@ -28,6 +28,7 @@ import (
 	jobparams "github.com/kubeflow/pipelines/backend/api/go_http_client/job_client/job_service"
 	pipelineparams "github.com/kubeflow/pipelines/backend/api/go_http_client/pipeline_client/pipeline_service"
 	runparams "github.com/kubeflow/pipelines/backend/api/go_http_client/run_client/run_service"
+	"github.com/kubeflow/pipelines/backend/api/go_http_client/run_model"
 	"github.com/kubeflow/pipelines/backend/src/common/client/api_server"
 	"github.com/pkg/errors"
 	"github.com/stretchr/testify/assert"
@@ -96,3 +97,14 @@ func DeleteAllJobs(client *api_server.JobClient, t *testing.T) {
 		assert.Nil(t, client.Delete(&jobparams.DeleteJobParams{ID: j.ID}))
 	}
 }
+
+func GetExperimentIDFromAPIResourceReferences(resourceRefs []*run_model.APIResourceReference) string {
+	experimentID := ""
+	for _, resourceRef := range resourceRefs {
+		if resourceRef.Key.Type == run_model.APIResourceTypeEXPERIMENT {
+			experimentID = resourceRef.Key.ID
+			break
+		}
+	}
+	return experimentID
+}