prysm-pulse/validator/db/kv/manage_test.go
terence tsao 2a997828a3
Audit best practice feedback - validator DB test folder (#6426)
* Refactor validators db

* Lint

* Lint

Co-authored-by: prylabs-bulldozer[bot] <58059840+prylabs-bulldozer[bot]@users.noreply.github.com>
2020-06-26 21:37:43 -05:00

296 lines
9.2 KiB
Go

package kv
import (
"bytes"
"context"
"encoding/hex"
"os"
"path/filepath"
"testing"
"github.com/pkg/errors"
"github.com/prysmaticlabs/go-bitfield"
slashpb "github.com/prysmaticlabs/prysm/proto/slashing"
"github.com/prysmaticlabs/prysm/shared/testutil"
bolt "go.etcd.io/bbolt"
)
type storeHistory struct {
Proposals map[[48]byte]bitfield.Bitlist
Attestations map[[48]byte]map[uint64]uint64
}
func TestMerge(t *testing.T) {
firstStorePubKeys := [][48]byte{{1}, {2}}
firstStore := setupDB(t, firstStorePubKeys)
secondStorePubKeys := [][48]byte{{3}, {4}}
secondStore := setupDB(t, secondStorePubKeys)
storeHistory1, err := prepareStore(firstStore, firstStorePubKeys)
if err != nil {
t.Fatal(err)
}
storeHistory2, err := prepareStore(secondStore, secondStorePubKeys)
if err != nil {
t.Fatal(err)
}
mergedProposals := make(map[[48]byte]bitfield.Bitlist)
for k, v := range storeHistory1.Proposals {
mergedProposals[k] = v
}
for k, v := range storeHistory2.Proposals {
mergedProposals[k] = v
}
mergedAttestations := make(map[[48]byte]map[uint64]uint64)
for k, v := range storeHistory1.Attestations {
mergedAttestations[k] = v
}
for k, v := range storeHistory2.Attestations {
mergedAttestations[k] = v
}
mergedStoreHistory := storeHistory{
Proposals: mergedProposals,
Attestations: mergedAttestations,
}
targetDirectory := testutil.TempDir() + "/target"
t.Cleanup(func() {
if err := os.RemoveAll(targetDirectory); err != nil {
t.Errorf("Could not remove target directory : %v", err)
}
})
err = Merge(context.Background(), []*Store{firstStore, secondStore}, targetDirectory)
if err != nil {
t.Fatalf("Merging failed: %v", err)
}
mergedStore, err := GetKVStore(targetDirectory)
if err != nil {
t.Fatalf("Retrieving the merged store failed: %v", err)
}
assertStore(
t,
mergedStore,
append(firstStorePubKeys, secondStorePubKeys[0], secondStorePubKeys[1]),
&mergedStoreHistory)
}
func TestSplit(t *testing.T) {
pubKey1 := [48]byte{1}
pubKey2 := [48]byte{2}
sourceStore := setupDB(t, [][48]byte{pubKey1, pubKey2})
storeHistory1, err := prepareStore(sourceStore, [][48]byte{pubKey1})
if err != nil {
t.Fatal(err)
}
storeHistory2, err := prepareStore(sourceStore, [][48]byte{pubKey2})
if err != nil {
t.Fatal(err)
}
targetDirectory := testutil.TempDir() + "/target"
t.Cleanup(func() {
if err := os.RemoveAll(targetDirectory); err != nil {
t.Errorf("Could not remove target directory: %v", err)
}
})
if err := Split(context.Background(), sourceStore, targetDirectory); err != nil {
t.Fatalf("Splitting failed: %v", err)
}
encodedKey1 := hex.EncodeToString(pubKey1[:])[:12]
keyStore1, err := GetKVStore(filepath.Join(targetDirectory, encodedKey1))
if err != nil {
t.Fatalf("Retrieving the store for public key %v failed: %v", encodedKey1, err)
}
if keyStore1 == nil {
t.Fatalf("No store created for public key %v", encodedKey1)
}
encodedKey2 := hex.EncodeToString(pubKey2[:])[:12]
keyStore2, err := GetKVStore(filepath.Join(targetDirectory, encodedKey2))
if err != nil {
t.Fatalf("Retrieving the store for public key %v failed: %v", encodedKey2, err)
}
if keyStore2 == nil {
t.Fatalf("No store created for public key %v", encodedKey2)
}
if err := keyStore1.view(func(tx *bolt.Tx) error {
otherKeyProposalsBucket := tx.Bucket(historicProposalsBucket).Bucket(pubKey2[:])
if otherKeyProposalsBucket != nil {
t.Fatalf("Store for public key %v contains proposals for another key", encodedKey2)
}
otherKeyAttestationsBucket := tx.Bucket(historicAttestationsBucket).Bucket(pubKey2[:])
if otherKeyAttestationsBucket != nil {
t.Fatalf("Store for public key %v contains attestations for another key", encodedKey2)
}
return nil
}); err != nil {
t.Fatalf("Failed to close store: %v", err)
}
if err := keyStore2.view(func(tx *bolt.Tx) error {
otherKeyProposalsBucket := tx.Bucket(historicProposalsBucket).Bucket(pubKey1[:])
if otherKeyProposalsBucket != nil {
t.Fatalf("Store for public key %v contains proposals for another key", encodedKey1)
}
otherKeyAttestationsBucket := tx.Bucket(historicAttestationsBucket).Bucket(pubKey1[:])
if otherKeyAttestationsBucket != nil {
t.Fatalf("Store for public key %v contains attestations for another key", encodedKey1)
}
return nil
}); err != nil {
t.Fatalf("Failed to close store: %v", err)
}
assertStore(t, keyStore1, [][48]byte{pubKey1}, storeHistory1)
assertStore(t, keyStore2, [][48]byte{pubKey2}, storeHistory2)
}
func TestSplit_AttestationsWithoutMatchingProposalsAreSplit(t *testing.T) {
pubKey1 := [48]byte{1}
pubKey2 := [48]byte{2}
sourceStore := setupDB(t, [][48]byte{pubKey1, pubKey2})
_, err := prepareStoreProposals(sourceStore, [][48]byte{pubKey1})
if err != nil {
t.Fatal(err)
}
attestationHistory, err := prepareStoreAttestations(sourceStore, [][48]byte{pubKey1, pubKey2})
if err != nil {
t.Fatal(err)
}
targetDirectory := testutil.TempDir() + "/target"
t.Cleanup(func() {
if err := os.RemoveAll(targetDirectory); err != nil {
t.Errorf("Could not remove target directory : %v", err)
}
})
if err := Split(context.Background(), sourceStore, targetDirectory); err != nil {
t.Fatalf("Splitting failed: %v", err)
}
encodedKey1 := hex.EncodeToString(pubKey1[:])[:12]
encodedKey2 := hex.EncodeToString(pubKey2[:])[:12]
attestationsOnlyKeyStore, err := GetKVStore(filepath.Join(targetDirectory, encodedKey2))
if err != nil {
t.Fatalf("Retrieving the store failed: %v", err)
}
if attestationsOnlyKeyStore == nil {
t.Fatalf("No store created for public key %v", encodedKey2)
}
if err := attestationsOnlyKeyStore.view(func(tx *bolt.Tx) error {
otherKeyProposalsBucket := tx.Bucket(historicProposalsBucket).Bucket(pubKey1[:])
if otherKeyProposalsBucket != nil {
t.Fatalf("Store for public key %v contains proposals for another key", encodedKey1)
}
otherKeyAttestationsBucket := tx.Bucket(historicAttestationsBucket).Bucket(pubKey1[:])
if otherKeyAttestationsBucket != nil {
t.Fatalf("Store for public key %v contains attestations for another key", encodedKey1)
}
return nil
}); err != nil {
t.Fatalf("Failed to retrieve attestations: %v", err)
}
splitAttestationsHistory, err :=
attestationsOnlyKeyStore.AttestationHistoryForPubKeys(context.Background(), [][48]byte{pubKey2})
if err != nil {
t.Fatalf("Retrieving attestation history failed for public key %v", encodedKey2)
}
if splitAttestationsHistory[pubKey2].TargetToSource[0] != attestationHistory[pubKey2][0] {
t.Fatalf(
"Attestations not merged correctly: expected %v vs received %v",
attestationHistory[pubKey2][0],
splitAttestationsHistory[pubKey2].TargetToSource[0])
}
}
func prepareStore(store *Store, pubKeys [][48]byte) (*storeHistory, error) {
proposals, err := prepareStoreProposals(store, pubKeys)
if err != nil {
return nil, err
}
attestations, err := prepareStoreAttestations(store, pubKeys)
if err != nil {
return nil, err
}
history := storeHistory{
Proposals: proposals,
Attestations: attestations,
}
return &history, nil
}
func prepareStoreProposals(store *Store, pubKeys [][48]byte) (map[[48]byte]bitfield.Bitlist, error) {
proposals := make(map[[48]byte]bitfield.Bitlist)
for i, key := range pubKeys {
proposalHistory := bitfield.Bitlist{byte(i), 0x00, 0x00, 0x00, 0x01}
if err := store.SaveProposalHistoryForEpoch(context.Background(), key[:], 0, proposalHistory); err != nil {
return nil, errors.Wrapf(err, "Saving proposal history failed")
}
proposals[key] = proposalHistory
}
return proposals, nil
}
func prepareStoreAttestations(store *Store, pubKeys [][48]byte) (map[[48]byte]map[uint64]uint64, error) {
storeAttestationHistory := make(map[[48]byte]*slashpb.AttestationHistory)
attestations := make(map[[48]byte]map[uint64]uint64)
for i, key := range pubKeys {
attestationHistoryMap := make(map[uint64]uint64)
attestationHistoryMap[0] = uint64(i)
attestationHistory := &slashpb.AttestationHistory{
TargetToSource: attestationHistoryMap,
LatestEpochWritten: 0,
}
storeAttestationHistory[key] = attestationHistory
attestations[key] = attestationHistoryMap
}
if err := store.SaveAttestationHistoryForPubKeys(context.Background(), storeAttestationHistory); err != nil {
return nil, errors.Wrapf(err, "Saving attestation history failed")
}
return attestations, nil
}
func assertStore(t *testing.T, store *Store, pubKeys [][48]byte, expectedHistory *storeHistory) {
for _, key := range pubKeys {
proposalHistory, err := store.ProposalHistoryForEpoch(
context.Background(), key[:], 0)
if err != nil {
t.Fatalf("Retrieving proposal history failed for public key %v", key)
}
expectedProposals := expectedHistory.Proposals[key]
if !bytes.Equal(proposalHistory, expectedProposals) {
t.Fatalf("Proposals are incorrect: expected %v vs received %v", expectedProposals, proposalHistory)
}
}
attestationHistory, err := store.AttestationHistoryForPubKeys(context.Background(), pubKeys)
if err != nil {
t.Fatalf("Retrieving attestation history failed")
}
for _, key := range pubKeys {
expectedAttestations := expectedHistory.Attestations[key]
if attestationHistory[key].TargetToSource[0] != expectedAttestations[0] {
t.Fatalf(
"Attestations are incorrect: expected %v vs received %v",
expectedAttestations[0],
attestationHistory[key].TargetToSource[0])
}
}
}