Return pubkey with AttestationHistoryForPubKey AttestationRecord resp…

…onse (#9135) * Return pubkey with AttestationHistoryForPubKey AttestationRecord response * Fix tests
prysmaticlabs · Jul 1, 2021 · 2df024a · 2df024a
1 parent ca7b312
commit 2df024a
Show file tree

Hide file tree

Showing 5 changed files with 11 additions and 8 deletions.
diff --git a/validator/db/kv/attester_protection.go b/validator/db/kv/attester_protection.go
@@ -114,6 +114,7 @@ func (s *Store) AttestationHistoryForPubKey(ctx context.Context, pubKey [48]byte
 			sourceEpoch := bytesutil.BytesToEpochBigEndian(sourceBytes)
 			for _, targetEpoch := range targetEpochs {
 				record := &AttestationRecord{
+					PubKey: pubKey,
 					Source: sourceEpoch,
 					Target: targetEpoch,
 				}

diff --git a/validator/slashing-protection/cli_import_export_test.go b/validator/slashing-protection/cli_import_export_test.go
@@ -46,7 +46,7 @@ func TestImportExportSlashingProtectionCli_RoundTrip(t *testing.T) {
 	// Create some mock slashing protection history. and JSON file
 	pubKeys, err := mocks.CreateRandomPubKeys(numValidators)
 	require.NoError(t, err)
-	attestingHistory, proposalHistory := mocks.MockAttestingAndProposalHistories(numValidators)
+	attestingHistory, proposalHistory := mocks.MockAttestingAndProposalHistories(pubKeys)
 	require.NoError(t, err)
 	mockJSON, err := mocks.MockSlashingProtectionJSON(pubKeys, attestingHistory, proposalHistory)
 	require.NoError(t, err)

diff --git a/validator/slashing-protection/local/standard-protection-format/import_test.go b/validator/slashing-protection/local/standard-protection-format/import_test.go
@@ -53,7 +53,7 @@ func TestStore_ImportInterchangeData_BadFormat_PreventsDBWrites(t *testing.T) {
 
 	// First we setup some mock attesting and proposal histories and create a mock
 	// standard slashing protection format JSON struct.
-	attestingHistory, proposalHistory := valtest.MockAttestingAndProposalHistories(numValidators)
+	attestingHistory, proposalHistory := valtest.MockAttestingAndProposalHistories(publicKeys)
 	standardProtectionFormat, err := valtest.MockSlashingProtectionJSON(publicKeys, attestingHistory, proposalHistory)
 	require.NoError(t, err)
 
@@ -111,7 +111,7 @@ func TestStore_ImportInterchangeData_OK(t *testing.T) {
 
 	// First we setup some mock attesting and proposal histories and create a mock
 	// standard slashing protection format JSON struct.
-	attestingHistory, proposalHistory := valtest.MockAttestingAndProposalHistories(numValidators)
+	attestingHistory, proposalHistory := valtest.MockAttestingAndProposalHistories(publicKeys)
 	standardProtectionFormat, err := valtest.MockSlashingProtectionJSON(publicKeys, attestingHistory, proposalHistory)
 	require.NoError(t, err)
 

diff --git a/validator/slashing-protection/local/standard-protection-format/round_trip_test.go b/validator/slashing-protection/local/standard-protection-format/round_trip_test.go
@@ -26,7 +26,7 @@ func TestImportExport_RoundTrip(t *testing.T) {
 
 	// First we setup some mock attesting and proposal histories and create a mock
 	// standard slashing protection format JSON struct.
-	attestingHistory, proposalHistory := slashtest.MockAttestingAndProposalHistories(numValidators)
+	attestingHistory, proposalHistory := slashtest.MockAttestingAndProposalHistories(publicKeys)
 	require.NoError(t, err)
 	wanted, err := slashtest.MockSlashingProtectionJSON(publicKeys, attestingHistory, proposalHistory)
 	require.NoError(t, err)
@@ -137,7 +137,7 @@ func TestImportInterchangeData_OK(t *testing.T) {
 
 	// First we setup some mock attesting and proposal histories and create a mock
 	// standard slashing protection format JSON struct.
-	attestingHistory, proposalHistory := slashtest.MockAttestingAndProposalHistories(numValidators)
+	attestingHistory, proposalHistory := slashtest.MockAttestingAndProposalHistories(publicKeys)
 	require.NoError(t, err)
 	standardProtectionFormat, err := slashtest.MockSlashingProtectionJSON(publicKeys, attestingHistory, proposalHistory)
 	require.NoError(t, err)
@@ -196,7 +196,7 @@ func TestImportInterchangeData_OK_SavesBlacklistedPublicKeys(t *testing.T) {
 
 	// First we setup some mock attesting and proposal histories and create a mock
 	// standard slashing protection format JSON struct.
-	attestingHistory, proposalHistory := slashtest.MockAttestingAndProposalHistories(numValidators)
+	attestingHistory, proposalHistory := slashtest.MockAttestingAndProposalHistories(publicKeys)
 	require.NoError(t, err)
 
 	standardProtectionFormat, err := slashtest.MockSlashingProtectionJSON(publicKeys, attestingHistory, proposalHistory)
@@ -284,7 +284,7 @@ func TestStore_ImportInterchangeData_BadFormat_PreventsDBWrites(t *testing.T) {
 
 	// First we setup some mock attesting and proposal histories and create a mock
 	// standard slashing protection format JSON struct.
-	attestingHistory, proposalHistory := slashtest.MockAttestingAndProposalHistories(numValidators)
+	attestingHistory, proposalHistory := slashtest.MockAttestingAndProposalHistories(publicKeys)
 	require.NoError(t, err)
 	standardProtectionFormat, err := slashtest.MockSlashingProtectionJSON(publicKeys, attestingHistory, proposalHistory)
 	require.NoError(t, err)

diff --git a/validator/testing/protection_history.go b/validator/testing/protection_history.go
@@ -51,8 +51,9 @@ func MockSlashingProtectionJSON(
 
 // MockAttestingAndProposalHistories given a number of validators, creates mock attesting
 // and proposing histories within WEAK_SUBJECTIVITY_PERIOD bounds.
-func MockAttestingAndProposalHistories(numValidators int) ([][]*kv.AttestationRecord, []kv.ProposalHistoryForPubkey) {
+func MockAttestingAndProposalHistories(pubkeys [][48]byte) ([][]*kv.AttestationRecord, []kv.ProposalHistoryForPubkey) {
 	// deduplicate and transform them into our internal format.
+	numValidators := len(pubkeys)
 	attData := make([][]*kv.AttestationRecord, numValidators)
 	proposalData := make([]kv.ProposalHistoryForPubkey, numValidators)
 	gen := rand.NewGenerator()
@@ -73,6 +74,7 @@ func MockAttestingAndProposalHistories(numValidators int) ([][]*kv.AttestationRe
 				Source:      i - 1,
 				Target:      i,
 				SigningRoot: signingRoot,
+				PubKey:      pubkeys[v],
 			})
 		}
 		for i := types.Epoch(1); i <= latestTarget; i++ {