fix flaky test

Signed-off-by: Charlie Le <[email protected]>
cortexproject · Oct 27, 2024 · 6a3b4b4 · 6a3b4b4
1 parent d711e1b
commit 6a3b4b4
Showing 1 changed file with 6 additions and 4 deletions.
diff --git a/pkg/ingester/ingester_test.go b/pkg/ingester/ingester_test.go
@@ -5315,12 +5315,14 @@ func Test_Ingester_ModeHandler(t *testing.T) {
 
 			require.Equal(t, testData.expectedResponse, response.Code)
 			require.Equal(t, testData.expectedState, i.lifecycler.GetState())
-
-			err = i.CheckReady(context.Background())
 			if testData.expectedIsReady {
-				require.NoError(t, err)
+				// Wait for instance to own tokens
+				test.Poll(t, 1*time.Second, nil, func() interface{} {
+					return i.CheckReady(context.Background())
+				})
+				require.NoError(t, i.CheckReady(context.Background()))
 			} else {
-				require.NotNil(t, err)
+				require.NotNil(t, i.CheckReady(context.Background()))
 			}
 		})
 	}