improve checks

segmentio · Oct 10, 2023 · c1ae689 · c1ae689
1 parent 3bb0bb6
commit c1ae689
Showing 1 changed file with 19 additions and 0 deletions.
diff --git a/pkg/admin/brokerclient_test.go b/pkg/admin/brokerclient_test.go
@@ -455,6 +455,25 @@ func TestBrokerClientAlterAssignments(t *testing.T) {
 		topicInfo, err = client.GetTopic(ctx, topicName, true)
 		require.NoError(t, err)
 
+		// Alter partition succeeded
+		if topicInfo.Partitions[2].Replicas[0] != 5 {
+			return fmt.Errorf("Assign partitions change not reflected yet")
+		}
+
+		for _, partition := range topicInfo.Partitions {
+			if len(partition.Replicas) != 2 {
+				return fmt.Errorf("Assign partitions change not reflected yet")
+			}
+		}
+
+		return nil
+	})
+
+	util.RetryUntil(t, 5*time.Second, func() error {
+		topicInfo, err = client.GetTopic(ctx, topicName, true)
+		require.NoError(t, err)
+
+		// ISR shrink completed
 		for _, partition := range topicInfo.Partitions {
 			if len(partition.Replicas) != 2 {
 				return fmt.Errorf("Assign partitions change not reflected yet")