workload test framework improvements: enhance privatekey access (#1798)

cli/spam.go

@@ -6,6 +6,7 @@ package cli
 import (
 	"context"
 
+	"github.com/ava-labs/hypersdk/auth"
 	"github.com/ava-labs/hypersdk/cli/prompt"
 	"github.com/ava-labs/hypersdk/consts"
 	"github.com/ava-labs/hypersdk/throughput"
@@ -43,8 +44,13 @@ func (h *Handler) BuildSpammer(sh throughput.SpamHelper, defaults bool) (*throug
 		return nil, err
 	}
 
+	authFactory, err := auth.GetFactory(key)
+	if err != nil {
+		return nil, err
+	}
+
 	if defaults {
-		sc := throughput.NewDefaultConfig(uris, key)
+		sc := throughput.NewDefaultConfig(uris, authFactory)
 		return throughput.NewSpammer(sc, sh)
 	}
 	// Collect parameters
@@ -83,7 +89,7 @@ func (h *Handler) BuildSpammer(sh throughput.SpamHelper, defaults bool) (*throug
 
 	sc := throughput.NewConfig(
 		uris,
-		key,
+		authFactory,
 		sZipf,
 		vZipf,
 		txsPerSecond,

docs/tutorials/morpheusvm/4_testing.md

@@ -286,33 +286,24 @@ func newDefaultKeys() []ed25519.PrivateKey {
 }
 ```
 
-Finally, we implement the network configuration required for our VM
+Finally, we initialize the workload.DefaultTestNetworkConfiguration required for our VM
 tests:
 
 ```go
-type NetworkConfiguration struct {
-	workload.DefaultTestNetworkConfiguration
-	keys []ed25519.PrivateKey
-}
-
-func (n *NetworkConfiguration) Keys() []ed25519.PrivateKey {
-	return n.keys
-}
 
 func NewTestNetworkConfig(minBlockGap time.Duration) (*NetworkConfiguration, error) {
 	keys := newDefaultKeys()
 	genesis := newGenesis(keys, minBlockGap)
 	genesisBytes, err := json.Marshal(genesis)
 	if err != nil {
-		return nil, err
+		return workload.DefaultTestNetworkConfiguration{}, err
 	}
-	return &NetworkConfiguration{
-		DefaultTestNetworkConfiguration: workload.NewDefaultTestNetworkConfiguration(
-			genesisBytes,
-			consts.Name,
-			vm.NewParser(genesis)),
-		keys: keys,
-	}, nil
+	return workload.NewDefaultTestNetworkConfiguration(
+		genesisBytes,
+		consts.Name,
+		vm.NewParser(genesis),
+		keys,
+	), nil
 }
 ```
 
@@ -394,8 +385,7 @@ function:
 	require.NoError(err)
 	toAddress := auth.NewED25519Address(other.PublicKey())
 
-	networkConfig := tn.Configuration().(*workload.NetworkConfiguration)
-	spendingKey := networkConfig.Keys()[0]
+	authFactory := tn.Configuration().AuthFactories()[0]
 ```
 
 Next, we'll create our test transaction. In short, we'll want to send a value of
@@ -406,7 +396,7 @@ Next, we'll create our test transaction. In short, we'll want to send a value of
 		To:    toAddress,
 		Value: 1,
 	}},
-		auth.NewED25519Factory(spendingKey),
+		authFactory,
 	)
 	require.NoError(err)
 ```
@@ -474,7 +464,7 @@ var _ = ginkgo.BeforeSuite(func() {
 
 	randomEd25519AuthFactory := auth.NewED25519Factory(randomEd25519Priv)
 
-	generator := workload.NewTxGenerator(testingNetworkConfig.Keys()[0])
+	generator := workload.NewTxGenerator(testingNetworkConfig.AuthFactories()[0])
 	// Setup imports the integration test coverage
 	integration.Setup(
 		vm.New,

docs/tutorials/morpheusvm/5_interlude.md

@@ -105,14 +105,10 @@ var _ = ginkgo.SynchronizedBeforeSuite(func() []byte {
 	expectedABI, err := abi.NewABI(vm.ActionParser.GetRegisteredTypes(), vm.OutputParser.GetRegisteredTypes())
 	require.NoError(err)
 
-	firstKey := testingNetworkConfig.Keys()[0]
-	generator := workload.NewTxGenerator(firstKey)
-	spamKey := &auth.PrivateKey{
-		Address: auth.NewED25519Address(firstKey.PublicKey()),
-		Bytes:   firstKey[:],
-	}
+	firstAuthFactory := testingNetworkConfig.AuthFactories()[0]
+	generator := workload.NewTxGenerator(firstAuthFactory)
 	tc := e2e.NewTestContext()
-	he2e.SetWorkload(testingNetworkConfig, generator, expectedABI, nil, spamKey)
+	he2e.SetWorkload(testingNetworkConfig, generator, expectedABI, nil, firstAuthFactory)
 
 	return fixture.NewTestEnvironment(tc, flagVars, owner, testingNetworkConfig, consts.ID).Marshal()
 }, func(envBytes []byte) {

examples/morpheusvm/tests/e2e/e2e_test.go

@@ -52,14 +52,10 @@ var _ = ginkgo.SynchronizedBeforeSuite(func() []byte {
 		KeyType: auth.ED25519Key,
 	}
 
-	firstKey := testingNetworkConfig.Keys()[0]
-	generator := workload.NewTxGenerator(firstKey)
-	spamKey := &auth.PrivateKey{
-		Address: auth.NewED25519Address(firstKey.PublicKey()),
-		Bytes:   firstKey[:],
-	}
+	firstAuthFactory := testingNetworkConfig.AuthFactories()[0]
+	generator := workload.NewTxGenerator(firstAuthFactory)
 	tc := e2e.NewTestContext()
-	he2e.SetWorkload(testingNetworkConfig, generator, expectedABI, &spamHelper, spamKey)
+	he2e.SetWorkload(testingNetworkConfig, generator, expectedABI, &spamHelper, firstAuthFactory)
 
 	return fixture.NewTestEnvironment(tc, flagVars, owner, testingNetworkConfig, consts.ID).Marshal()
 }, func(envBytes []byte) {

examples/morpheusvm/tests/integration/integration_test.go

@@ -35,7 +35,7 @@ var _ = ginkgo.BeforeSuite(func() {
 
 	randomEd25519AuthFactory := auth.NewED25519Factory(randomEd25519Priv)
 
-	generator := workload.NewTxGenerator(testingNetworkConfig.Keys()[0])
+	generator := workload.NewTxGenerator(testingNetworkConfig.AuthFactories()[0])
 	// Setup imports the integration test coverage
 	integration.Setup(
 		vm.New,

examples/morpheusvm/tests/transfer.go

@@ -13,7 +13,6 @@ import (
 	"github.com/ava-labs/hypersdk/chain"
 	"github.com/ava-labs/hypersdk/crypto/ed25519"
 	"github.com/ava-labs/hypersdk/examples/morpheusvm/actions"
-	"github.com/ava-labs/hypersdk/examples/morpheusvm/tests/workload"
 	"github.com/ava-labs/hypersdk/tests/registry"
 
 	tworkload "github.com/ava-labs/hypersdk/tests/workload"
@@ -31,14 +30,12 @@ var _ = registry.Register(TestsRegistry, "Transfer Transaction", func(t ginkgo.F
 	require.NoError(err)
 	toAddress := auth.NewED25519Address(other.PublicKey())
 
-	networkConfig := tn.Configuration().(*workload.NetworkConfiguration)
-	spendingKey := networkConfig.Keys()[0]
-
+	authFactory := tn.Configuration().AuthFactories()[0]
 	tx, err := tn.GenerateTx(context.Background(), []chain.Action{&actions.Transfer{
 		To:    toAddress,
 		Value: 1,
 	}},
-		auth.NewED25519Factory(spendingKey),
+		authFactory,
 	)
 	require.NoError(err)
 

examples/morpheusvm/tests/workload/generator.go

@@ -27,12 +27,12 @@ var _ workload.TxGenerator = (*TxGenerator)(nil)
 const txCheckInterval = 100 * time.Millisecond
 
 type TxGenerator struct {
-	factory *auth.ED25519Factory
+	factory chain.AuthFactory
 }
 
-func NewTxGenerator(key ed25519.PrivateKey) *TxGenerator {
+func NewTxGenerator(authFactory chain.AuthFactory) *TxGenerator {
 	return &TxGenerator{
-		factory: auth.NewED25519Factory(key),
+		factory: authFactory,
 	}
 }
 

examples/morpheusvm/tests/workload/genesis.go

@@ -11,6 +11,7 @@ import (
 	"github.com/ava-labs/avalanchego/ids"
 
 	"github.com/ava-labs/hypersdk/auth"
+	"github.com/ava-labs/hypersdk/chain"
 	"github.com/ava-labs/hypersdk/codec"
 	"github.com/ava-labs/hypersdk/crypto/ed25519"
 	"github.com/ava-labs/hypersdk/examples/morpheusvm/consts"
@@ -25,20 +26,18 @@ const (
 	InitialBalance uint64 = 10_000_000_000_000
 )
 
-var _ workload.TestNetworkConfiguration = &NetworkConfiguration{}
-
 // hardcoded initial set of ed25519 keys. Each will be initialized with InitialBalance
 var ed25519HexKeys = []string{
 	"323b1d8f4eed5f0da9da93071b034f2dce9d2d22692c172f3cb252a64ddfafd01b057de320297c29ad0c1f589ea216869cf1938d88c9fbd70d6748323dbf2fa7", //nolint:lll
 	"8a7be2e0c9a2d09ac2861c34326d6fe5a461d920ba9c2b345ae28e603d517df148735063f8d5d8ba79ea4668358943e5c80bc09e9b2b9a15b5b15db6c1862e88", //nolint:lll
 }
 
-func newGenesis(keys []ed25519.PrivateKey, minBlockGap time.Duration) *genesis.DefaultGenesis {
+func newGenesis(authFactories []chain.AuthFactory, minBlockGap time.Duration) *genesis.DefaultGenesis {
 	// allocate the initial balance to the addresses
-	customAllocs := make([]*genesis.CustomAllocation, 0, len(keys))
-	for _, key := range keys {
+	customAllocs := make([]*genesis.CustomAllocation, 0, len(authFactories))
+	for _, authFactory := range authFactories {
 		customAllocs = append(customAllocs, &genesis.CustomAllocation{
-			Address: auth.NewED25519Address(key.PublicKey()),
+			Address: authFactory.Address(),
 			Balance: InitialBalance,
 		})
 	}
@@ -59,40 +58,29 @@ func newGenesis(keys []ed25519.PrivateKey, minBlockGap time.Duration) *genesis.D
 	return genesis
 }
 
-func newDefaultKeys() []ed25519.PrivateKey {
-	testKeys := make([]ed25519.PrivateKey, len(ed25519HexKeys))
+func newDefaultAuthFactories() []chain.AuthFactory {
+	authFactories := make([]chain.AuthFactory, len(ed25519HexKeys))
 	for i, keyHex := range ed25519HexKeys {
 		bytes, err := codec.LoadHex(keyHex, ed25519.PrivateKeyLen)
 		if err != nil {
 			panic(err)
 		}
-		testKeys[i] = ed25519.PrivateKey(bytes)
+		authFactories[i] = auth.NewED25519Factory(ed25519.PrivateKey(bytes))
 	}
-
-	return testKeys
-}
-
-type NetworkConfiguration struct {
-	workload.DefaultTestNetworkConfiguration
-	keys []ed25519.PrivateKey
-}
-
-func (n *NetworkConfiguration) Keys() []ed25519.PrivateKey {
-	return n.keys
+	return authFactories
 }
 
-func NewTestNetworkConfig(minBlockGap time.Duration) (*NetworkConfiguration, error) {
-	keys := newDefaultKeys()
+func NewTestNetworkConfig(minBlockGap time.Duration) (workload.DefaultTestNetworkConfiguration, error) {
+	keys := newDefaultAuthFactories()
 	genesis := newGenesis(keys, minBlockGap)
 	genesisBytes, err := json.Marshal(genesis)
 	if err != nil {
-		return nil, err
+		return workload.DefaultTestNetworkConfiguration{}, err
 	}
-	return &NetworkConfiguration{
-		DefaultTestNetworkConfiguration: workload.NewDefaultTestNetworkConfiguration(
-			genesisBytes,
-			consts.Name,
-			vm.NewParser(genesis)),
-		keys: keys,
-	}, nil
+	return workload.NewDefaultTestNetworkConfiguration(
+		genesisBytes,
+		consts.Name,
+		vm.NewParser(genesis),
+		keys,
+	), nil
 }

tests/e2e/e2e.go

@@ -18,7 +18,7 @@ import (
 	"github.com/ava-labs/hypersdk/abi"
 	"github.com/ava-labs/hypersdk/api/jsonrpc"
 	"github.com/ava-labs/hypersdk/api/state"
-	"github.com/ava-labs/hypersdk/auth"
+	"github.com/ava-labs/hypersdk/chain"
 	"github.com/ava-labs/hypersdk/tests/registry"
 	"github.com/ava-labs/hypersdk/tests/workload"
 	"github.com/ava-labs/hypersdk/throughput"
@@ -31,11 +31,11 @@ var (
 	networkConfig workload.TestNetworkConfiguration
 	txWorkload    workload.TxWorkload
 	expectedABI   abi.ABI
-	spamKey       *auth.PrivateKey
+	spamKey       chain.AuthFactory
 	spamHelper    throughput.SpamHelper
 )
 
-func SetWorkload(networkConfigImpl workload.TestNetworkConfiguration, generator workload.TxGenerator, abi abi.ABI, sh throughput.SpamHelper, key *auth.PrivateKey) {
+func SetWorkload(networkConfigImpl workload.TestNetworkConfiguration, generator workload.TxGenerator, abi abi.ABI, sh throughput.SpamHelper, key chain.AuthFactory) {
 	networkConfig = networkConfigImpl
 	txWorkload = workload.TxWorkload{
 		Generator: generator,

tests/e2e/network.go

@@ -21,6 +21,8 @@ import (
 var (
 	ErrUnableToConfirmTx = errors.New("unable to confirm transaction")
 	ErrInvalidURI        = errors.New("invalid uri")
+
+	_ workload.TestNetwork = (*Network)(nil)
 )
 
 const (

tests/integration/network.go

@@ -20,6 +20,8 @@ var (
 	ErrUnableToConfirmTx = errors.New("unable to confirm transaction")
 	ErrInvalidURI        = errors.New("invalid uri")
 	ErrTxNotFound        = errors.New("tx not found")
+
+	_ workload.TestNetwork = (*Network)(nil)
 )
 
 type Network struct {

tests/workload/network.go

@@ -23,14 +23,16 @@ type TestNetworkConfiguration interface {
 	GenesisBytes() []byte
 	Name() string
 	Parser() chain.Parser
+	AuthFactories() []chain.AuthFactory
 }
 
 // DefaultTestNetworkConfiguration struct is the common test configuration that a test framework would need to provide
 // in order to deploy a network. A test would typically embed this as part of it's network configuration structure.
 type DefaultTestNetworkConfiguration struct {
-	genesisBytes []byte
-	name         string
-	parser       chain.Parser
+	genesisBytes  []byte
+	name          string
+	parser        chain.Parser
+	authFactories []chain.AuthFactory
 }
 
 func (d DefaultTestNetworkConfiguration) GenesisBytes() []byte {
@@ -45,11 +47,16 @@ func (d DefaultTestNetworkConfiguration) Parser() chain.Parser {
 	return d.parser
 }
 
+func (d DefaultTestNetworkConfiguration) AuthFactories() []chain.AuthFactory {
+	return d.authFactories
+}
+
 // NewDefaultTestNetworkConfiguration creates a new DefaultTestNetworkConfiguration object.
-func NewDefaultTestNetworkConfiguration(genesisBytes []byte, name string, parser chain.Parser) DefaultTestNetworkConfiguration {
+func NewDefaultTestNetworkConfiguration(genesisBytes []byte, name string, parser chain.Parser, authFactories []chain.AuthFactory) DefaultTestNetworkConfiguration {
 	return DefaultTestNetworkConfiguration{
-		genesisBytes: genesisBytes,
-		name:         name,
-		parser:       parser,
+		genesisBytes:  genesisBytes,
+		name:          name,
+		parser:        parser,
+		authFactories: authFactories,
 	}
 }

throughput/config.go

@@ -3,11 +3,13 @@
 
 package throughput
 
-import "github.com/ava-labs/hypersdk/auth"
+import (
+	"github.com/ava-labs/hypersdk/chain"
+)
 
 type Config struct {
 	uris             []string
-	key              *auth.PrivateKey
+	authFactory      chain.AuthFactory
 	sZipf            float64
 	vZipf            float64
 	txsPerSecond     int
@@ -19,11 +21,11 @@ type Config struct {
 
 func NewDefaultConfig(
 	uris []string,
-	key *auth.PrivateKey,
+	authFactory chain.AuthFactory,
 ) *Config {
 	return &Config{
 		uris:             uris,
-		key:              key,
+		authFactory:      authFactory,
 		sZipf:            1.01,
 		vZipf:            2.7,
 		txsPerSecond:     500,
@@ -36,7 +38,7 @@ func NewDefaultConfig(
 
 func NewConfig(
 	uris []string,
-	key *auth.PrivateKey,
+	key chain.AuthFactory,
 	sZipf float64,
 	vZipf float64,
 	txsPerSecond int,