From 0476e21b097003068f56c62c35d59b87e6ba7a30 Mon Sep 17 00:00:00 2001
From: Nikita Kryuchkov <nkryuchkov10@gmail.com>
Date: Mon, 14 Oct 2024 21:56:04 -0300
Subject: [PATCH] fix TestSetupValidatorsExporter

---
 operator/validator/controller_test.go  | 25 +++++++++++++++---------
 operator/validator/mocks/controller.go | 27 +++++++++++++-------------
 2 files changed, 30 insertions(+), 22 deletions(-)

diff --git a/operator/validator/controller_test.go b/operator/validator/controller_test.go
index eca857d05d..6f534108a8 100644
--- a/operator/validator/controller_test.go
+++ b/operator/validator/controller_test.go
@@ -6,6 +6,7 @@ import (
 	"encoding/base64"
 	"encoding/hex"
 	"fmt"
+	"slices"
 	"sync"
 	"testing"
 	"time"
@@ -29,6 +30,7 @@ import (
 	ibftstorage "github.com/ssvlabs/ssv/ibft/storage"
 	"github.com/ssvlabs/ssv/logging"
 	"github.com/ssvlabs/ssv/network"
+	"github.com/ssvlabs/ssv/network/commons"
 	"github.com/ssvlabs/ssv/networkconfig"
 	operatordatastore "github.com/ssvlabs/ssv/operator/datastore"
 	"github.com/ssvlabs/ssv/operator/keys"
@@ -192,15 +194,14 @@ func TestSetupValidatorsExporter(t *testing.T) {
 	testCases := []struct {
 		name                       string
 		shareStorageListResponse   []*types.SSVShare
-		expectMetadataFetch        bool
 		syncHighestDecidedResponse error
 		getValidatorDataResponse   error
 	}{
-		{"no shares of non committee", nil, false, nil, nil},
-		{"set up non committee validators", sharesWithMetadata, false, nil, nil},
-		{"set up non committee validators without metadata", sharesWithoutMetadata, false, nil, nil},
-		{"fail to sync highest decided", sharesWithMetadata, false, errors.New("failed to sync highest decided"), nil},
-		{"fail to update validators metadata", sharesWithMetadata, false, nil, errors.New("could not update all validators")},
+		{"no shares of non committee", nil, nil, nil},
+		{"set up non committee validators", sharesWithMetadata, nil, nil},
+		{"set up non committee validators without metadata", sharesWithoutMetadata, nil, nil},
+		{"fail to sync highest decided", sharesWithMetadata, errors.New("failed to sync highest decided"), nil},
+		{"fail to update validators metadata", sharesWithMetadata, nil, errors.New("could not update all validators")},
 	}
 
 	for _, tc := range testCases {
@@ -210,6 +211,14 @@ func TestSetupValidatorsExporter(t *testing.T) {
 			defer ctrl.Finish()
 			mockValidatorsMap := validators.New(context.TODO())
 
+			var subnets []byte
+			for _, share := range sharesWithMetadata {
+				subnets = append(subnets, byte(commons.ValidatorSubnet(hex.EncodeToString(share.ValidatorPubKey[:]))))
+			}
+			slices.Sort(subnets)
+
+			network.EXPECT().ActiveSubnets().Return(subnets).AnyTimes()
+
 			if tc.shareStorageListResponse == nil {
 				sharesStorage.EXPECT().List(gomock.Any(), gomock.Any()).Return(tc.shareStorageListResponse).Times(1)
 			} else {
@@ -229,9 +238,7 @@ func TestSetupValidatorsExporter(t *testing.T) {
 						}
 					}
 				}).AnyTimes()
-				if tc.expectMetadataFetch {
-					bc.EXPECT().GetValidatorData(gomock.Any()).Return(bcResponse, tc.getValidatorDataResponse).Times(1)
-				}
+				bc.EXPECT().GetValidatorData(gomock.Any()).Return(bcResponse, tc.getValidatorDataResponse).AnyTimes()
 				bc.EXPECT().GetBeaconNetwork().Return(networkconfig.Mainnet.Beacon.GetBeaconNetwork()).AnyTimes()
 				sharesStorage.EXPECT().UpdateValidatorsMetadata(gomock.Any()).Return(nil).AnyTimes()
 				recipientStorage.EXPECT().GetRecipientData(gomock.Any(), gomock.Any()).Return(recipientData, true, nil).AnyTimes()
diff --git a/operator/validator/mocks/controller.go b/operator/validator/mocks/controller.go
index f66ca22518..3c61c36242 100644
--- a/operator/validator/mocks/controller.go
+++ b/operator/validator/mocks/controller.go
@@ -5,22 +5,23 @@
 package mocks
 
 import (
-	reflect "reflect"
+	"reflect"
 
-	phase0 "github.com/attestantio/go-eth2-client/spec/phase0"
-	common "github.com/ethereum/go-ethereum/common"
-	types "github.com/ssvlabs/ssv-spec-pre-cc/types"
+	"github.com/attestantio/go-eth2-client/spec/phase0"
+	"github.com/ethereum/go-ethereum/common"
+	"github.com/ssvlabs/ssv-spec-pre-cc/types"
 	types0 "github.com/ssvlabs/ssv-spec/types"
-	network "github.com/ssvlabs/ssv/network"
-	records "github.com/ssvlabs/ssv/network/records"
-	duties "github.com/ssvlabs/ssv/operator/duties"
-	validators "github.com/ssvlabs/ssv/operator/validators"
-	beacon "github.com/ssvlabs/ssv/protocol/v2/blockchain/beacon"
+	"go.uber.org/mock/gomock"
+	"go.uber.org/zap"
+
+	"github.com/ssvlabs/ssv/network"
+	"github.com/ssvlabs/ssv/network/records"
+	"github.com/ssvlabs/ssv/operator/duties"
+	"github.com/ssvlabs/ssv/operator/validators"
+	"github.com/ssvlabs/ssv/protocol/v2/blockchain/beacon"
 	types1 "github.com/ssvlabs/ssv/protocol/v2/types"
-	storage "github.com/ssvlabs/ssv/registry/storage"
-	basedb "github.com/ssvlabs/ssv/storage/basedb"
-	gomock "go.uber.org/mock/gomock"
-	zap "go.uber.org/zap"
+	"github.com/ssvlabs/ssv/registry/storage"
+	"github.com/ssvlabs/ssv/storage/basedb"
 )
 
 // MockController is a mock of Controller interface.