Skip to content

Commit 66e8886

Browse files
committed
Make test key generator usable from benchmarks
1 parent 8e4ab70 commit 66e8886

File tree

4 files changed

+40
-41
lines changed

4 files changed

+40
-41
lines changed

rolling-shutter/medley/testkeygen/keygenerator.go

Lines changed: 12 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -14,21 +14,20 @@ import (
1414
// TestKeyGenerator is a helper tool to generate secret and public eon and epoch keys and key
1515
// shares. It will generate a new eon key every eonInterval epochs.
1616
type TestKeyGenerator struct {
17-
t *testing.T
17+
tb testing.TB
1818
eonInterval uint64
1919
eonKeyGen map[uint64]*EonKeys
2020
NumKeypers uint64
2121
Threshold uint64
2222
}
2323

24-
func NewTestKeyGenerator(t *testing.T, numKeypers uint64, threshold uint64, infiniteInterval bool) *TestKeyGenerator {
25-
t.Helper()
26-
eonInterval := 100
24+
func NewTestKeyGenerator(tb testing.TB, numKeypers uint64, threshold uint64, infiniteInterval bool) *TestKeyGenerator {
25+
tb.Helper()
2726
if infiniteInterval {
2827
eonInterval = 0 // 0 stands for infinity
2928
}
3029
return &TestKeyGenerator{
31-
t: t,
30+
tb: tb,
3231
eonInterval: uint64(eonInterval),
3332
eonKeyGen: make(map[uint64]*EonKeys),
3433
NumKeypers: numKeypers,
@@ -47,7 +46,7 @@ func (tkg *TestKeyGenerator) getEonIndex(identityPreimage identitypreimage.Ident
4746
}
4847

4948
func (tkg *TestKeyGenerator) EonKeysForEpoch(identityPreimage identitypreimage.IdentityPreimage) *EonKeys {
50-
tkg.t.Helper()
49+
tkg.tb.Helper()
5150
var err error
5251
eonIndex := tkg.getEonIndex(identityPreimage)
5352
res, ok := tkg.eonKeyGen[eonIndex]
@@ -57,7 +56,7 @@ func (tkg *TestKeyGenerator) EonKeysForEpoch(identityPreimage identitypreimage.I
5756
tkg.NumKeypers,
5857
tkg.Threshold,
5958
)
60-
assert.NilError(tkg.t, err)
59+
assert.NilError(tkg.tb, err)
6160
tkg.eonKeyGen[eonIndex] = res
6261
}
6362
return res
@@ -66,32 +65,32 @@ func (tkg *TestKeyGenerator) EonKeysForEpoch(identityPreimage identitypreimage.I
6665
func (tkg *TestKeyGenerator) EonPublicKeyShare(identityPreimage identitypreimage.IdentityPreimage,
6766
keyperIndex uint64,
6867
) *shcrypto.EonPublicKeyShare {
69-
tkg.t.Helper()
68+
tkg.tb.Helper()
7069
return tkg.EonKeysForEpoch(identityPreimage).keyperShares[keyperIndex].eonPublicKeyShare
7170
}
7271

7372
func (tkg *TestKeyGenerator) EonPublicKey(identityPreimage identitypreimage.IdentityPreimage) *shcrypto.EonPublicKey {
74-
tkg.t.Helper()
73+
tkg.tb.Helper()
7574
return tkg.EonKeysForEpoch(identityPreimage).publicKey
7675
}
7776

7877
func (tkg *TestKeyGenerator) EonSecretKeyShare(identityPreimage identitypreimage.IdentityPreimage,
7978
keyperIndex uint64,
8079
) *shcrypto.EonSecretKeyShare {
81-
tkg.t.Helper()
80+
tkg.tb.Helper()
8281
return tkg.EonKeysForEpoch(identityPreimage).keyperShares[keyperIndex].eonSecretKeyShare
8382
}
8483

8584
func (tkg *TestKeyGenerator) EpochSecretKeyShare(identityPreimage identitypreimage.IdentityPreimage,
8685
keyperIndex uint64,
8786
) *shcrypto.EpochSecretKeyShare {
88-
tkg.t.Helper()
87+
tkg.tb.Helper()
8988
return tkg.EonKeysForEpoch(identityPreimage).keyperShares[keyperIndex].ComputeEpochSecretKeyShare(identityPreimage)
9089
}
9190

9291
func (tkg *TestKeyGenerator) EpochSecretKey(identityPreimage identitypreimage.IdentityPreimage) *shcrypto.EpochSecretKey {
93-
tkg.t.Helper()
92+
tkg.tb.Helper()
9493
epochSecretKey, err := tkg.EonKeysForEpoch(identityPreimage).EpochSecretKey(identityPreimage)
95-
assert.NilError(tkg.t, err)
94+
assert.NilError(tkg.tb, err)
9695
return epochSecretKey
9796
}

rolling-shutter/medley/testsetup/database.go

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -42,47 +42,47 @@ var testDBSuffix = "-test"
4242
// newDBPoolTeardown connects to a test db specified an environment variable and clears it from all
4343
// schemas we might have created. It returns the db connection pool and a close function. Call the
4444
// close function at the end of the test to reset the db again and close the connection.
45-
func newDBPoolTeardown(ctx context.Context, t *testing.T) (*pgxpool.Pool, func()) {
46-
t.Helper()
45+
func newDBPoolTeardown(ctx context.Context, tb testing.TB) (*pgxpool.Pool, func()) {
46+
tb.Helper()
4747

4848
testDBURL, exists := os.LookupEnv(testDBURLVar)
4949
if !exists {
50-
t.Skipf("no test db specified, please set %s", testDBURLVar)
50+
tb.Skipf("no test db specified, please set %s", testDBURLVar)
5151
}
5252

5353
dbpool, err := pgxpool.Connect(ctx, testDBURL)
5454
if err != nil {
55-
t.Fatalf("failed to connect to test db: %v", err)
55+
tb.Fatalf("failed to connect to test db: %v", err)
5656
}
5757

5858
closedb := func() {
5959
_, err = dbpool.Exec(ctx, dropEverything)
6060
dbpool.Close() // close db no matter if dropping failed
6161
if err != nil {
62-
t.Fatalf("failed to reset test db: %v", err)
62+
tb.Fatalf("failed to reset test db: %v", err)
6363
}
6464
}
6565

6666
// drop db contents
6767
_, err = dbpool.Exec(ctx, dropEverything)
6868
if err != nil {
6969
dbpool.Close()
70-
t.Fatalf("failed to reset test db: %v", err)
70+
tb.Fatalf("failed to reset test db: %v", err)
7171
}
7272

7373
return dbpool, closedb
7474
}
7575

76-
func NewTestDBPool(ctx context.Context, t *testing.T, definition db.Definition) (*pgxpool.Pool, func()) {
77-
t.Helper()
76+
func NewTestDBPool(ctx context.Context, tb testing.TB, definition db.Definition) (*pgxpool.Pool, func()) {
77+
tb.Helper()
7878

79-
dbpool, closedb := newDBPoolTeardown(ctx, t)
79+
dbpool, closedb := newDBPoolTeardown(ctx, tb)
8080

8181
err := db.InitDB(ctx, dbpool, definition.Name()+testDBSuffix, definition)
8282
if err != nil {
8383
log.Error().Err(err).Str("db-definition", definition.Name()).Msg("Initializing DB failed")
8484
closedb()
85-
t.Fatalf("failed to initialize '%s' db", definition.Name())
85+
tb.Fatalf("failed to initialize '%s' db", definition.Name())
8686
}
8787
return dbpool, closedb
8888
}

rolling-shutter/medley/testsetup/eon.go

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -31,15 +31,15 @@ type TestConfig interface {
3131

3232
func InitializeEon(
3333
ctx context.Context,
34-
t *testing.T,
34+
tb testing.TB,
3535
dbpool *pgxpool.Pool,
3636
config TestConfig,
3737
keyperIndex uint64,
3838
) *testkeygen.TestKeyGenerator {
39-
t.Helper()
39+
tb.Helper()
4040

4141
err := dbpool.BeginFunc(db.WrapContext(ctx, database.Definition.Validate))
42-
assert.NilError(t, err)
42+
assert.NilError(tb, err)
4343

4444
keyperDB := database.New(dbpool)
4545
keypers := []string{
@@ -51,16 +51,16 @@ func InitializeEon(
5151
collatorKey := config.GetCollatorKey()
5252
if collatorKey != nil {
5353
err := dbpool.BeginFunc(db.WrapContext(ctx, chainobsdb.Definition.Validate))
54-
assert.NilError(t, err)
54+
assert.NilError(tb, err)
5555
chdb := chainobsdb.New(dbpool)
5656
err = chdb.InsertChainCollator(ctx, chainobsdb.InsertChainCollatorParams{
5757
ActivationBlockNumber: 0,
5858
Collator: shdb.EncodeAddress(ethcrypto.PubkeyToAddress(config.GetCollatorKey().PublicKey)),
5959
})
60-
assert.NilError(t, err)
60+
assert.NilError(tb, err)
6161
}
6262

63-
tkg := testkeygen.NewTestKeyGenerator(t, 3, 2, false)
63+
tkg := testkeygen.NewTestKeyGenerator(tb, 3, 2, false)
6464
publicKeyShares := []*shcrypto.EonPublicKeyShare{}
6565
identityPreimage := identitypreimage.BigToIdentityPreimage(common.Big0)
6666
for i := uint64(0); i < tkg.NumKeypers; i++ {
@@ -77,29 +77,29 @@ func InitializeEon(
7777
PublicKeyShares: publicKeyShares,
7878
}
7979
dkgResultEncoded, err := shdb.EncodePureDKGResult(&dkgResult)
80-
assert.NilError(t, err)
80+
assert.NilError(tb, err)
8181

8282
err = keyperDB.InsertBatchConfig(ctx, database.InsertBatchConfigParams{
8383
KeyperConfigIndex: 1,
8484
Height: 0,
8585
Keypers: keypers,
8686
Threshold: int32(tkg.Threshold),
8787
})
88-
assert.NilError(t, err)
88+
assert.NilError(tb, err)
8989
err = keyperDB.InsertEon(ctx, database.InsertEonParams{
9090
Eon: int64(config.GetEon()),
9191
Height: 0,
9292
ActivationBlockNumber: 0,
9393
KeyperConfigIndex: 1,
9494
})
95-
assert.NilError(t, err)
95+
assert.NilError(tb, err)
9696
err = keyperDB.InsertDKGResult(ctx, database.InsertDKGResultParams{
9797
Eon: int64(config.GetEon()),
9898
Success: true,
9999
Error: sql.NullString{},
100100
PureResult: dkgResultEncoded,
101101
})
102-
assert.NilError(t, err)
102+
assert.NilError(tb, err)
103103

104104
return tkg
105105
}

rolling-shutter/p2p/p2ptest/p2ptest.go

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -16,35 +16,35 @@ import (
1616
// MustValidateMessageResult calls the handlers ValidateMessage method and ensures it returns the
1717
// expected result.
1818
func MustValidateMessageResult(
19-
t *testing.T,
19+
tb testing.TB,
2020
expectedResult pubsub.ValidationResult,
2121
handler p2p.MessageHandler,
2222
ctx context.Context, //nolint:revive
2323
msg p2pmsg.Message,
2424
) {
25-
t.Helper()
25+
tb.Helper()
2626
validationResult, err := handler.ValidateMessage(ctx, msg)
2727
accepted := validationResult == pubsub.ValidationAccept
2828
log.Debug().
2929
Interface("msg", msg).
3030
Int("result", int(validationResult)).
3131
Int("expected", int(expectedResult)).Err(err).Msg("ValidateMessage")
3232
if accepted {
33-
assert.NilError(t, err, "validation returned error")
33+
assert.NilError(tb, err, "validation returned error")
3434
}
35-
assert.Equal(t, expectedResult, validationResult, "validation did not validate with expected result ")
35+
assert.Equal(tb, expectedResult, validationResult, "validation did not validate with expected result ")
3636
}
3737

3838
// MustHandleMessage makes sure the handler validates and handles the given message without errors.
3939
func MustHandleMessage(
40-
t *testing.T,
40+
tb testing.TB,
4141
handler p2p.MessageHandler,
4242
ctx context.Context, //nolint:revive
4343
msg p2pmsg.Message,
4444
) []p2pmsg.Message {
45-
t.Helper()
46-
MustValidateMessageResult(t, pubsub.ValidationAccept, handler, ctx, msg)
45+
tb.Helper()
46+
MustValidateMessageResult(tb, pubsub.ValidationAccept, handler, ctx, msg)
4747
msgs, err := handler.HandleMessage(ctx, msg)
48-
assert.NilError(t, err)
48+
assert.NilError(tb, err)
4949
return msgs
5050
}

0 commit comments

Comments
 (0)