cap execution of test

SeldonIO · Nov 1, 2024 · dfb3472 · dfb3472
1 parent b7a70b6
commit dfb3472
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 4 deletions.
diff --git a/scheduler/pkg/agent/server_test.go b/scheduler/pkg/agent/server_test.go
@@ -1054,9 +1054,13 @@ func TestSubscribe(t *testing.T) {
 				}(a.id)
 			}
 
-			for len(server.agents) != test.expectedAgentsCount {
+			maxCount := 10
+			count := 0
+			for len(server.agents) != test.expectedAgentsCount && count < maxCount {
 				time.Sleep(100 * time.Millisecond)
+				count++
 			}
+			g.Expect(len(server.agents)).To(Equal(test.expectedAgentsCount))
 
 			for idx, s := range streams {
 				go func(idx int, s *grpc.ClientConn) {
@@ -1066,9 +1070,12 @@ func TestSubscribe(t *testing.T) {
 				}(idx, s)
 			}
 
-			for len(server.agents) != test.expectedAgentsCountAfterClose {
+			count = 0
+			for len(server.agents) != test.expectedAgentsCountAfterClose && count < maxCount {
 				time.Sleep(100 * time.Millisecond)
+				count++
 			}
+			g.Expect(len(server.agents)).To(Equal(test.expectedAgentsCountAfterClose))
 
 			server.StopAgentStreams()
 		})

diff --git a/scheduler/pkg/kafka/dataflow/server_test.go b/scheduler/pkg/kafka/dataflow/server_test.go
@@ -639,6 +639,7 @@ func TestPipelineRebalance(t *testing.T) {
 }
 
 func TestPipelineSubscribe(t *testing.T) {
+	g := NewGomegaWithT(t)
 	type ag struct {
 		id      uint32
 		doClose bool
@@ -739,9 +740,13 @@ func TestPipelineSubscribe(t *testing.T) {
 				}(a.id)
 			}
 
-			for len(s.streams) != test.expectedAgentsCount {
+			maxCount := 10
+			count := 0
+			for len(s.streams) != test.expectedAgentsCount && count < maxCount {
 				time.Sleep(100 * time.Millisecond)
+				count++
 			}
+			g.Expect(len(s.streams)).To(Equal(test.expectedAgentsCount))
 
 			for idx, s := range streams {
 				go func(idx int, s *grpc.ClientConn) {
@@ -751,9 +756,12 @@ func TestPipelineSubscribe(t *testing.T) {
 				}(idx, s)
 			}
 
-			for len(s.streams) != test.expectedAgentsCountAfterClose {
+			count = 0
+			for len(s.streams) != test.expectedAgentsCountAfterClose && count < maxCount {
 				time.Sleep(100 * time.Millisecond)
+				count++
 			}
+			g.Expect(len(s.streams)).To(Equal(test.expectedAgentsCountAfterClose))
 
 			s.StopSendPipelineEvents()
 		})