Deflake TestWaitForActivation_AccountsChanged (#12282)

Co-authored-by: terencechain <terence@prysmaticlabs.com>
prysmaticlabs · Apr 17, 2023 · 898cb0b · 898cb0b
1 parent 6944d22
commit 898cb0b
Show file tree

Hide file tree

Showing 4 changed files with 14 additions and 5 deletions.
diff --git a/go.mod b/go.mod
@@ -211,6 +211,7 @@ require (
 	github.com/russross/blackfriday/v2 v2.1.0 // indirect
 	github.com/shirou/gopsutil v3.21.11+incompatible // indirect
 	github.com/spaolacci/murmur3 v1.1.0 // indirect
+	github.com/stretchr/objx v0.5.0 // indirect
 	github.com/syndtr/goleveldb v1.0.1-0.20220721030215-126854af5e6d // indirect
 	github.com/tklauser/go-sysconf v0.3.11 // indirect
 	github.com/tklauser/numcpus v0.6.0 // indirect

diff --git a/go.sum b/go.sum
@@ -1186,6 +1186,7 @@ github.com/streadway/handy v0.0.0-20190108123426-d5acb3125c2a/go.mod h1:qNTQ5P5J
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
 github.com/stretchr/objx v0.1.1/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
 github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw=
+github.com/stretchr/objx v0.5.0 h1:1zr/of2m5FGMsad5YfcqgdqdWrIhu+EBEJRhR1U7z/c=
 github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo=
 github.com/stretchr/testify v1.2.0/go.mod h1:a8OnRcib4nhh0OaRAV+Yts87kKdq0PP7pXfy6kDkUVs=
 github.com/stretchr/testify v1.2.2/go.mod h1:a8OnRcib4nhh0OaRAV+Yts87kKdq0PP7pXfy6kDkUVs=

diff --git a/validator/client/BUILD.bazel b/validator/client/BUILD.bazel
@@ -162,6 +162,7 @@ go_test(
         "@com_github_prysmaticlabs_go_bitfield//:go_default_library",
         "@com_github_sirupsen_logrus//:go_default_library",
         "@com_github_sirupsen_logrus//hooks/test:go_default_library",
+        "@com_github_stretchr_testify//mock:go_default_library",
         "@com_github_tyler_smith_go_bip39//:go_default_library",
         "@com_github_wealdtech_go_eth2_util//:go_default_library",
         "@in_gopkg_d4l3k_messagediff_v1//:go_default_library",

diff --git a/validator/client/wait_for_activation_test.go b/validator/client/wait_for_activation_test.go
@@ -1,6 +1,7 @@
 package client
 
 import (
+	"bytes"
 	"context"
 	"fmt"
 	"testing"
@@ -19,6 +20,7 @@ import (
 	"github.com/prysmaticlabs/prysm/v4/validator/keymanager/derived"
 	constant "github.com/prysmaticlabs/prysm/v4/validator/testing"
 	logTest "github.com/sirupsen/logrus/hooks/test"
+	mock2 "github.com/stretchr/testify/mock"
 	"github.com/tyler-smith/go-bip39"
 	util "github.com/wealdtech/go-eth2-util"
 )
@@ -251,8 +253,6 @@ func TestWaitForActivation_RefetchKeys(t *testing.T) {
 
 // Regression test for a scenario where you start with an inactive key and then import an active key.
 func TestWaitForActivation_AccountsChanged(t *testing.T) {
-	t.Skip("Flakey test. Skipping until we can figure out how to test this properly")
-
 	hook := logTest.NewGlobal()
 	ctrl := gomock.NewController(t)
 	defer ctrl.Finish()
@@ -299,9 +299,15 @@ func TestWaitForActivation_AccountsChanged(t *testing.T) {
 		activeClientStream := mock.NewMockBeaconNodeValidator_WaitForActivationClient(ctrl)
 		validatorClient.EXPECT().WaitForActivation(
 			gomock.Any(),
-			&ethpb.ValidatorActivationRequest{
-				PublicKeys: [][]byte{inactivePubKey[:], activePubKey[:]},
-			},
+			mock2.MatchedBy(func(req *ethpb.ValidatorActivationRequest) bool {
+				found := 0
+				for _, pk := range req.PublicKeys {
+					if bytes.Equal(pk, activePubKey[:]) || bytes.Equal(pk, inactivePubKey[:]) {
+						found++
+					}
+				}
+				return found == 2
+			}),
 		).Return(activeClientStream, nil)
 		activeClientStream.EXPECT().Recv().Return(
 			activeResp,