Fix flaky snapshot test.

criteo-forks · Nov 6, 2023 · 5cb614e · 5cb614e
1 parent 28b1469
commit 5cb614e
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/agent/snapshot_endpoint_test.go b/agent/snapshot_endpoint_test.go
@@ -74,6 +74,7 @@ func TestSnapshot_Options(t *testing.T) {
 		t.Run(method, func(t *testing.T) {
 			a := NewTestAgent(t, TestACLConfig())
 			defer a.Shutdown()
+			testrpc.WaitForLeader(t, a.RPC, "dc1")
 
 			body := bytes.NewBuffer(nil)
 			req, _ := http.NewRequest(method, "/v1/snapshot", body)
@@ -88,6 +89,7 @@ func TestSnapshot_Options(t *testing.T) {
 		t.Run(method, func(t *testing.T) {
 			a := NewTestAgent(t, TestACLConfig())
 			defer a.Shutdown()
+			testrpc.WaitForLeader(t, a.RPC, "dc1")
 
 			body := bytes.NewBuffer(nil)
 			req, _ := http.NewRequest(method, "/v1/snapshot?dc=nope", body)
@@ -101,6 +103,7 @@ func TestSnapshot_Options(t *testing.T) {
 		t.Run(method, func(t *testing.T) {
 			a := NewTestAgent(t, TestACLConfig())
 			defer a.Shutdown()
+			testrpc.WaitForLeader(t, a.RPC, "dc1")
 
 			body := bytes.NewBuffer(nil)
 			req, _ := http.NewRequest(method, "/v1/snapshot?stale", body)