prysm-pulse/validator/db/kv/manage_test.go
Victor Farazdagi 2f11e55869
Use t.TempDir() in tests (#7769)
* use t.TempDir()

* remove redundant delete

* simplify setupDB()

* simplify db/testing/setup_db

* fix tests
2020-11-10 22:45:17 +00:00

212 lines
8.1 KiB
Go

package kv
import (
"context"
"encoding/hex"
"path/filepath"
"testing"
"github.com/pkg/errors"
slashpb "github.com/prysmaticlabs/prysm/proto/slashing"
"github.com/prysmaticlabs/prysm/shared/bytesutil"
"github.com/prysmaticlabs/prysm/shared/testutil/require"
bolt "go.etcd.io/bbolt"
)
type storeHistory struct {
Proposals map[[48]byte][]byte
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)
require.NoError(t, err)
storeHistory2, err := prepareStore(secondStore, secondStorePubKeys)
require.NoError(t, err)
mergedProposals := make(map[[48]byte][]byte)
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 := t.TempDir() + "/target"
err = Merge(context.Background(), []*Store{firstStore, secondStore}, targetDirectory)
require.NoError(t, err, "Merging failed")
mergedStore, err := GetKVStore(targetDirectory)
require.NoError(t, err, "Retrieving the merged store failed")
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})
require.NoError(t, err)
storeHistory2, err := prepareStore(sourceStore, [][48]byte{pubKey2})
require.NoError(t, err)
targetDirectory := t.TempDir() + "/target"
require.NoError(t, Split(context.Background(), sourceStore, targetDirectory), "Splitting failed")
encodedKey1 := hex.EncodeToString(pubKey1[:])[:12]
keyStore1, err := GetKVStore(filepath.Join(targetDirectory, encodedKey1))
require.NoError(t, err, "Retrieving the store for public key %v failed", encodedKey1)
require.NotNil(t, keyStore1, "No store created for public key %v", encodedKey1)
encodedKey2 := hex.EncodeToString(pubKey2[:])[:12]
keyStore2, err := GetKVStore(filepath.Join(targetDirectory, encodedKey2))
require.NoError(t, err, "Retrieving the store for public key %v failed", encodedKey2)
require.NotNil(t, keyStore2, "No store created for public key %v", encodedKey2)
err = keyStore1.view(func(tx *bolt.Tx) error {
otherKeyProposalsBucket := tx.Bucket(newhistoricProposalsBucket).Bucket(pubKey2[:])
require.Equal(t, (*bolt.Bucket)(nil), otherKeyProposalsBucket, "Store for public key %v contains proposals for another key", encodedKey2)
otherKeyAttestationsBucket := tx.Bucket(historicAttestationsBucket).Bucket(pubKey2[:])
require.Equal(t, (*bolt.Bucket)(nil), otherKeyAttestationsBucket, "Store for public key %v contains attestations for another key", encodedKey2)
return nil
})
require.NoError(t, err)
err = keyStore2.view(func(tx *bolt.Tx) error {
otherKeyProposalsBucket := tx.Bucket(newhistoricProposalsBucket).Bucket(pubKey1[:])
require.Equal(t, (*bolt.Bucket)(nil), otherKeyProposalsBucket, "Store for public key %v contains proposals for another key", encodedKey1)
otherKeyAttestationsBucket := tx.Bucket(historicAttestationsBucket).Bucket(pubKey1[:])
require.Equal(t, (*bolt.Bucket)(nil), otherKeyAttestationsBucket, "Store for public key %v contains attestations for another key", encodedKey1)
return nil
})
require.NoError(t, 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})
require.NoError(t, err)
attestationHistory, err := prepareStoreAttestations(sourceStore, [][48]byte{pubKey1, pubKey2})
require.NoError(t, err)
targetDirectory := t.TempDir() + "/target"
require.NoError(t, Split(context.Background(), sourceStore, targetDirectory), "Splitting failed")
encodedKey1 := hex.EncodeToString(pubKey1[:])[:12]
encodedKey2 := hex.EncodeToString(pubKey2[:])[:12]
attestationsOnlyKeyStore, err := GetKVStore(filepath.Join(targetDirectory, encodedKey2))
require.NoError(t, err, "Retrieving the store failed")
require.NotNil(t, attestationsOnlyKeyStore, "No store created for public key %v", encodedKey2)
err = attestationsOnlyKeyStore.view(func(tx *bolt.Tx) error {
otherKeyProposalsBucket := tx.Bucket(newhistoricProposalsBucket).Bucket(pubKey1[:])
require.Equal(t, (*bolt.Bucket)(nil), otherKeyProposalsBucket, "Store for public key %v contains proposals for another key", encodedKey1)
otherKeyAttestationsBucket := tx.Bucket(historicAttestationsBucket).Bucket(pubKey1[:])
require.Equal(t, (*bolt.Bucket)(nil), otherKeyAttestationsBucket, "Store for public key %v contains attestations for another key", encodedKey1)
return nil
})
require.NoError(t, err)
splitAttestationsHistory, err :=
attestationsOnlyKeyStore.AttestationHistoryForPubKeys(context.Background(), [][48]byte{pubKey2})
require.NoError(t, err, "Retrieving attestation history failed for public key %v", encodedKey2)
require.Equal(t, attestationHistory[pubKey2][0], splitAttestationsHistory[pubKey2].TargetToSource[0], "Attestations not merged correctly")
}
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][]byte, error) {
proposals := make(map[[48]byte][]byte)
for i, key := range pubKeys {
signingRoot := bytesutil.PadTo([]byte{byte(i)}, 32)
if err := store.SaveProposalHistoryForSlot(context.Background(), key[:], 0, signingRoot); err != nil {
return nil, errors.Wrapf(err, "Saving proposal history failed")
}
proposals[key] = signingRoot
}
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.ProposalHistoryForSlot(context.Background(), key[:], 0)
require.NoError(t, err, "Retrieving proposal history failed for public key %v", key)
expectedProposals := expectedHistory.Proposals[key]
require.DeepEqual(t, expectedProposals, proposalHistory, "Proposals are incorrect")
}
attestationHistory, err := store.AttestationHistoryForPubKeys(context.Background(), pubKeys)
require.NoError(t, err, "Retrieving attestation history failed")
for _, key := range pubKeys {
expectedAttestations := expectedHistory.Attestations[key]
require.Equal(t, expectedAttestations[0], attestationHistory[key].TargetToSource[0], "Attestations are incorrect")
}
}