diff --git a/tests/integration/channel_balance/channel_balance_test.go b/tests/integration/channel_balance/channel_balance_test.go
index 85a30b9b51..d69da10db2 100644
--- a/tests/integration/channel_balance/channel_balance_test.go
+++ b/tests/integration/channel_balance/channel_balance_test.go
@@ -113,7 +113,7 @@ func (s *ChannelBalanceSuite) flushCollections(collections []string) {
 			return info.GetCollectionID() == collID
 		})
 		lo.ForEach(collSegs, func(info *datapb.SegmentInfo, _ int) {
-			s.Require().Equal(commonpb.SegmentState_Flushed, info.GetState())
+			s.Require().Contains([]commonpb.SegmentState{commonpb.SegmentState_Flushed, commonpb.SegmentState_Flushing}, info.GetState())
 		})
 	}
 	log.Info("=========================Data flush done=========================")