cgu: added WaitUntilClusterInState and unit tests

This PR adds a few methods to allow waiting on the batch remediation progress for individual clusters on the CGU. Unit tests for these new methods are included.

Author: klaskosk

pkg/cgu/cgu.go

@@ -396,6 +396,71 @@ func (builder *CguBuilder) WaitUntilComplete(timeout time.Duration) (*CguBuilder
 	return builder.WaitForCondition(conditionComplete, timeout)
 }
 
+// WaitUntilClusterInState waits the specified timeout for a cluster in the CGU to be in the specified state.
+func (builder *CguBuilder) WaitUntilClusterInState(cluster, state string, timeout time.Duration) (*CguBuilder, error) {
+	if valid, err := builder.validate(); !valid {
+		return nil, err
+	}
+
+	if cluster == "" {
+		glog.V(100).Info("Cluster name cannot be empty")
+
+		return nil, fmt.Errorf("cluster name cannot be empty")
+	}
+
+	if state == "" {
+		glog.V(100).Info("State cannot be empty")
+
+		return nil, fmt.Errorf("state cannot be empty")
+	}
+
+	glog.V(100).Infof(
+		"Waiting until cluster %s on CGU %s in namespace %s is in state %s",
+		cluster, builder.Definition.Name, builder.Definition.Namespace, state)
+
+	if !builder.Exists() {
+		return nil, fmt.Errorf(
+			"cgu object %s does not exist in namespace %s", builder.Definition.Name, builder.Definition.Namespace)
+	}
+
+	var err error
+	err = wait.PollUntilContextTimeout(
+		context.TODO(), 3*time.Second, timeout, true, func(ctx context.Context) (bool, error) {
+			builder.Object, err = builder.apiClient.RanV1alpha1().ClusterGroupUpgrades(builder.Definition.Namespace).
+				Get(context.TODO(), builder.Definition.Name, metav1.GetOptions{})
+			if err != nil {
+				return false, nil
+			}
+
+			status, ok := builder.Object.Status.Status.CurrentBatchRemediationProgress[cluster]
+			if !ok {
+				glog.V(100).Infof(
+					"cluster %s not found in batch remediation progress for cgu %s in namespace %s",
+					cluster, builder.Definition.Name, builder.Definition.Namespace)
+
+				return false, nil
+			}
+
+			return status.State == state, nil
+		})
+
+	if err != nil {
+		return nil, err
+	}
+
+	return builder, nil
+}
+
+// WaitUntilClusterComplete waits the specified timeout for a cluster in the CGU to complete remidation.
+func (builder *CguBuilder) WaitUntilClusterComplete(cluster string, timeout time.Duration) (*CguBuilder, error) {
+	return builder.WaitUntilClusterInState(cluster, v1alpha1.Completed, timeout)
+}
+
+// WaitUntilClusterInProgress waits the specified timeout for a cluster in the CGU to start remidation.
+func (builder *CguBuilder) WaitUntilClusterInProgress(cluster string, timeout time.Duration) (*CguBuilder, error) {
+	return builder.WaitUntilClusterInState(cluster, v1alpha1.InProgress, timeout)
+}
+
 // WaitUntilBackupStarts waits the specified timeout for the backup to start.
 func (builder *CguBuilder) WaitUntilBackupStarts(timeout time.Duration) (*CguBuilder, error) {
 	if valid, err := builder.validate(); !valid {

pkg/cgu/cgu_test.go

@@ -17,6 +17,8 @@ var (
 	defaultCguName           = "cgu-test"
 	defaultCguNsName         = "test-ns"
 	defaultCguMaxConcurrency = 1
+	defaultCguClusterName    = "test-cluster"
+	defaultCguClusterState   = v1alpha1.NotStarted
 	defaultCguCondition      = conditionComplete
 )
 
@@ -551,6 +553,168 @@ func TestCguWaitUntilComplete(t *testing.T) {
 	}
 }
 
+func TestCguWaitUntilClusterInState(t *testing.T) {
+	testCases := []struct {
+		cluster       string
+		state         string
+		exists        bool
+		inState       bool
+		valid         bool
+		expectedError error
+	}{
+		{
+			cluster:       defaultCguClusterName,
+			state:         defaultCguClusterState,
+			exists:        true,
+			inState:       true,
+			valid:         true,
+			expectedError: nil,
+		},
+		{
+			cluster:       "",
+			state:         defaultCguClusterState,
+			exists:        true,
+			inState:       true,
+			valid:         true,
+			expectedError: fmt.Errorf("cluster name cannot be empty"),
+		},
+		{
+			cluster:       defaultCguClusterName,
+			state:         "",
+			exists:        true,
+			inState:       true,
+			valid:         true,
+			expectedError: fmt.Errorf("state cannot be empty"),
+		},
+		{
+			cluster:       defaultCguClusterName,
+			state:         defaultCguClusterState,
+			exists:        false,
+			inState:       true,
+			valid:         true,
+			expectedError: fmt.Errorf("cgu object %s does not exist in namespace %s", defaultCguName, defaultCguNsName),
+		},
+		{
+			cluster:       defaultCguClusterName,
+			state:         defaultCguClusterState,
+			exists:        true,
+			inState:       false,
+			valid:         true,
+			expectedError: context.DeadlineExceeded,
+		},
+		{
+			cluster:       defaultCguClusterName,
+			state:         defaultCguClusterState,
+			exists:        true,
+			inState:       true,
+			valid:         false,
+			expectedError: fmt.Errorf("CGU 'nsname' cannot be empty"),
+		},
+	}
+
+	for _, testCase := range testCases {
+		var (
+			runtimeObjects []runtime.Object
+			cguBuilder     *CguBuilder
+		)
+
+		if testCase.exists {
+			cgu := buildDummyCgu(defaultCguName, defaultCguNsName, defaultCguMaxConcurrency)
+
+			if testCase.inState {
+				cgu.Status.Status.CurrentBatchRemediationProgress = map[string]*v1alpha1.ClusterRemediationProgress{
+					testCase.cluster: {State: testCase.state},
+				}
+			}
+
+			runtimeObjects = append(runtimeObjects, cgu)
+		}
+
+		testSettings := clients.GetTestClients(clients.TestClientParams{
+			K8sMockObjects: runtimeObjects,
+		})
+
+		if testCase.valid {
+			cguBuilder = buildValidCguTestBuilder(testSettings)
+		} else {
+			cguBuilder = buildInvalidCguTestBuilder(testSettings)
+		}
+
+		_, err := cguBuilder.WaitUntilClusterInState(testCase.cluster, testCase.state, time.Second)
+		assert.Equal(t, testCase.expectedError, err)
+	}
+}
+
+func TestCguWaitUntilClusterComplete(t *testing.T) {
+	testCases := []struct {
+		complete      bool
+		expectedError error
+	}{
+		{
+			complete:      true,
+			expectedError: nil,
+		},
+		{
+			complete:      false,
+			expectedError: context.DeadlineExceeded,
+		},
+	}
+
+	for _, testCase := range testCases {
+		cgu := buildDummyCgu(defaultCguName, defaultCguNsName, defaultCguMaxConcurrency)
+
+		if testCase.complete {
+			cgu.Status.Status.CurrentBatchRemediationProgress = map[string]*v1alpha1.ClusterRemediationProgress{
+				defaultCguClusterName: {State: v1alpha1.Completed},
+			}
+		}
+
+		testSettings := clients.GetTestClients(clients.TestClientParams{
+			K8sMockObjects: []runtime.Object{cgu},
+		})
+
+		cguBuilder := buildValidCguTestBuilder(testSettings)
+		_, err := cguBuilder.WaitUntilClusterComplete(defaultCguClusterName, time.Second)
+
+		assert.Equal(t, testCase.expectedError, err)
+	}
+}
+
+func TestCguWaitUntilClusterInProgress(t *testing.T) {
+	testCases := []struct {
+		inProgress    bool
+		expectedError error
+	}{
+		{
+			inProgress:    true,
+			expectedError: nil,
+		},
+		{
+			inProgress:    false,
+			expectedError: context.DeadlineExceeded,
+		},
+	}
+
+	for _, testCase := range testCases {
+		cgu := buildDummyCgu(defaultCguName, defaultCguNsName, defaultCguMaxConcurrency)
+
+		if testCase.inProgress {
+			cgu.Status.Status.CurrentBatchRemediationProgress = map[string]*v1alpha1.ClusterRemediationProgress{
+				defaultCguClusterName: {State: v1alpha1.InProgress},
+			}
+		}
+
+		testSettings := clients.GetTestClients(clients.TestClientParams{
+			K8sMockObjects: []runtime.Object{cgu},
+		})
+
+		cguBuilder := buildValidCguTestBuilder(testSettings)
+		_, err := cguBuilder.WaitUntilClusterInProgress(defaultCguClusterName, time.Second)
+
+		assert.Equal(t, testCase.expectedError, err)
+	}
+}
+
 func TestWaitUntilBackupStarts(t *testing.T) {
 	cguObject := buildDummyCgu(defaultCguName, defaultCguNsName, defaultCguMaxConcurrency)
 	cguObject.Status.Backup = &v1alpha1.BackupStatus{}