prysm-pulse/slasher/db/testing/setup_db.go
Raul Jordan 546196a6fa
Other Package Godocs for Prysm (#5681)
* e2e docs
* slasher docs
* Merge branch 'other-package-godocs' of github.com:prysmaticlabs/prysm into other-package-godocs
* all validator package comments
* Merge branch 'master' into other-package-godocs
* completed all other packages
* Merge branch 'master' into other-package-godocs
* Merge refs/heads/master into other-package-godocs
2020-04-29 21:32:39 +00:00

64 lines
1.9 KiB
Go

// Package testing defines useful helper functions for unit tests with
// the slasher database.
package testing
import (
"crypto/rand"
"fmt"
"math/big"
"os"
"path"
"testing"
"github.com/prysmaticlabs/prysm/shared/testutil"
slasherDB "github.com/prysmaticlabs/prysm/slasher/db"
"github.com/prysmaticlabs/prysm/slasher/db/kv"
)
// SetupSlasherDB instantiates and returns a SlasherDB instance.
func SetupSlasherDB(t testing.TB, spanCacheEnabled bool) *kv.Store {
randPath, err := rand.Int(rand.Reader, big.NewInt(1000000))
if err != nil {
t.Fatalf("Could not generate random file path: %v", err)
}
p := path.Join(testutil.TempDir(), fmt.Sprintf("/%d", randPath))
if err := os.RemoveAll(p); err != nil {
t.Fatalf("Failed to remove directory: %v", err)
}
cfg := &kv.Config{}
db, err := slasherDB.NewDB(p, cfg)
db.EnableSpanCache(spanCacheEnabled)
if err != nil {
t.Fatalf("Failed to instantiate DB: %v", err)
}
return db
}
// SetupSlasherDBDiffCacheSize instantiates and returns a SlasherDB instance with non default cache size.
func SetupSlasherDBDiffCacheSize(t testing.TB, cacheItems int64, maxCacheSize int64) *kv.Store {
randPath, err := rand.Int(rand.Reader, big.NewInt(1000000))
if err != nil {
t.Fatalf("Could not generate random file path: %v", err)
}
p := path.Join(testutil.TempDir(), fmt.Sprintf("/%d", randPath))
if err := os.RemoveAll(p); err != nil {
t.Fatalf("Failed to remove directory: %v", err)
}
cfg := &kv.Config{}
newDB, err := slasherDB.NewDB(p, cfg)
if err != nil {
t.Fatalf("Failed to instantiate DB: %v", err)
}
return newDB
}
// TeardownSlasherDB cleans up a test SlasherDB instance.
func TeardownSlasherDB(t testing.TB, db *kv.Store) {
if err := db.Close(); err != nil {
t.Fatalf("Failed to close database: %v", err)
}
if err := os.RemoveAll(db.DatabasePath()); err != nil {
t.Fatalf("Failed to remove directory: %v", err)
}
}