Refactor Workload

examples/morpheusvm/tests/e2e/e2e_test.go

@@ -6,6 +6,7 @@ package e2e_test
 import (
 	"encoding/json"
 	"testing"
+	"time"
 
 	"github.com/ava-labs/avalanchego/tests/fixture/e2e"
 	"github.com/stretchr/testify/require"
@@ -38,12 +39,10 @@ func init() {
 var _ = ginkgo.SynchronizedBeforeSuite(func() []byte {
 	require := require.New(ginkgo.GinkgoT())
 
-	gen, workloadFactory, spamKey, err := workload.New(100 /* minBlockGap: 100ms */)
+	keys := workload.NewDefaultKeys()
+	genesis := workload.NewGenesis(keys, 100*time.Millisecond)
+	genesisBytes, err := json.Marshal(genesis)
 	require.NoError(err)
-
-	genesisBytes, err := json.Marshal(gen)
-	require.NoError(err)
-
 	expectedABI, err := abi.NewABI(vm.ActionParser.GetRegisteredTypes(), vm.OutputParser.GetRegisteredTypes())
 	require.NoError(err)
 
@@ -56,8 +55,13 @@ var _ = ginkgo.SynchronizedBeforeSuite(func() []byte {
 		KeyType: auth.ED25519Key,
 	}
 
+	generator := workload.NewTxGenerator(keys[0])
+	spamKey := &auth.PrivateKey{
+		Address: auth.NewED25519Address(keys[0].PublicKey()),
+		Bytes:   keys[0][:],
+	}
 	tc := e2e.NewTestContext()
-	he2e.SetWorkload(consts.Name, workloadFactory, expectedABI, parser, &spamHelper, spamKey)
+	he2e.SetWorkload(consts.Name, generator, expectedABI, parser, &spamHelper, spamKey)
 
 	return fixture.NewTestEnvironment(tc, flagVars, owner, consts.Name, consts.ID, genesisBytes).Marshal()
 }, func(envBytes []byte) {

examples/morpheusvm/tests/integration/integration_test.go

@@ -11,11 +11,11 @@ import (
 
 	"github.com/ava-labs/hypersdk/auth"
 	"github.com/ava-labs/hypersdk/crypto/ed25519"
+	"github.com/ava-labs/hypersdk/examples/morpheusvm/tests/workload"
 	"github.com/ava-labs/hypersdk/examples/morpheusvm/vm"
 	"github.com/ava-labs/hypersdk/tests/integration"
 
 	lconsts "github.com/ava-labs/hypersdk/examples/morpheusvm/consts"
-	morpheusWorkload "github.com/ava-labs/hypersdk/examples/morpheusvm/tests/workload"
 	ginkgo "github.com/onsi/ginkgo/v2"
 )
 
@@ -25,9 +25,9 @@ func TestIntegration(t *testing.T) {
 
 var _ = ginkgo.BeforeSuite(func() {
 	require := require.New(ginkgo.GinkgoT())
-	genesis, workloadFactory, _, err := morpheusWorkload.New(0 /* minBlockGap: 0ms */)
-	require.NoError(err)
 
+	keys := workload.NewDefaultKeys()
+	genesis := workload.NewGenesis(keys, 0)
 	genesisBytes, err := json.Marshal(genesis)
 	require.NoError(err)
 
@@ -36,13 +36,14 @@ var _ = ginkgo.BeforeSuite(func() {
 
 	randomEd25519AuthFactory := auth.NewED25519Factory(randomEd25519Priv)
 
+	generator := workload.NewTxGenerator(keys[0])
 	// Setup imports the integration test coverage
 	integration.Setup(
 		vm.New,
 		genesisBytes,
 		lconsts.ID,
 		vm.CreateParser,
-		workloadFactory,
+		generator,
 		randomEd25519AuthFactory,
 	)
 })

examples/morpheusvm/tests/workload/generator.go

@@ -0,0 +1,94 @@
+// Copyright (C) 2024, Ava Labs, Inc. All rights reserved.
+// See the file LICENSE for licensing terms.
+
+package workload
+
+import (
+	"context"
+	"time"
+
+	"github.com/ava-labs/avalanchego/ids"
+	"github.com/stretchr/testify/require"
+
+	"github.com/ava-labs/hypersdk/api/indexer"
+	"github.com/ava-labs/hypersdk/api/jsonrpc"
+	"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/actions"
+	"github.com/ava-labs/hypersdk/examples/morpheusvm/consts"
+	"github.com/ava-labs/hypersdk/examples/morpheusvm/vm"
+	"github.com/ava-labs/hypersdk/tests/workload"
+)
+
+var _ workload.TxGenerator = (*TxGenerator)(nil)
+
+const txCheckInterval = 100 * time.Millisecond
+
+type TxGenerator struct {
+	factory *auth.ED25519Factory
+}
+
+func NewTxGenerator(key ed25519.PrivateKey) *TxGenerator {
+	return &TxGenerator{
+		factory: auth.NewED25519Factory(key),
+	}
+}
+
+func (g *TxGenerator) GenerateTx(ctx context.Context, uri string) (*chain.Transaction, workload.TxAssertion, error) {
+	// TODO: no need to generate the clients every tx
+	cli := jsonrpc.NewJSONRPCClient(uri)
+	lcli := vm.NewJSONRPCClient(uri)
+
+	to, err := ed25519.GeneratePrivateKey()
+	if err != nil {
+		return nil, nil, err
+	}
+
+	toAddress := auth.NewED25519Address(to.PublicKey())
+	parser, err := lcli.Parser(ctx)
+	if err != nil {
+		return nil, nil, err
+	}
+	_, tx, _, err := cli.GenerateTransaction(
+		ctx,
+		parser,
+		[]chain.Action{&actions.Transfer{
+			To:    toAddress,
+			Value: 1,
+		}},
+		g.factory,
+	)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	return tx, func(ctx context.Context, require *require.Assertions, uri string) {
+		confirmTx(ctx, require, uri, tx.ID(), toAddress, 1)
+	}, nil
+}
+
+func confirmTx(ctx context.Context, require *require.Assertions, uri string, txID ids.ID, receiverAddr codec.Address, receiverExpectedBalance uint64) {
+	indexerCli := indexer.NewClient(uri)
+	success, _, err := indexerCli.WaitForTransaction(ctx, txCheckInterval, txID)
+	require.NoError(err)
+	require.True(success)
+	lcli := vm.NewJSONRPCClient(uri)
+	balance, err := lcli.Balance(ctx, receiverAddr)
+	require.NoError(err)
+	require.Equal(receiverExpectedBalance, balance)
+	txRes, _, err := indexerCli.GetTx(ctx, txID)
+	require.NoError(err)
+	// TODO: perform exact expected fee, units check, and output check
+	require.NotZero(txRes.Fee)
+	require.Len(txRes.Outputs, 1)
+	transferOutputBytes := []byte(txRes.Outputs[0])
+	require.Equal(consts.TransferID, transferOutputBytes[0])
+	reader := codec.NewReader(transferOutputBytes, len(transferOutputBytes))
+	transferOutputTyped, err := vm.OutputParser.Unmarshal(reader)
+	require.NoError(err)
+	transferOutput, ok := transferOutputTyped.(*actions.TransferResult)
+	require.True(ok)
+	require.Equal(receiverExpectedBalance, transferOutput.ReceiverBalance)
+}

examples/morpheusvm/tests/workload/genesis.go

@@ -0,0 +1,62 @@
+// Copyright (C) 2024, Ava Labs, Inc. All rights reserved.
+// See the file LICENSE for licensing terms.
+
+package workload
+
+import (
+	"math"
+	"time"
+
+	"github.com/ava-labs/hypersdk/auth"
+	"github.com/ava-labs/hypersdk/codec"
+	"github.com/ava-labs/hypersdk/crypto/ed25519"
+	"github.com/ava-labs/hypersdk/fees"
+	"github.com/ava-labs/hypersdk/genesis"
+)
+
+const (
+	// default initial balance for each address
+	InitialBalance uint64 = 10_000_000_000_000
+)
+
+// 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 {
+	// allocate the initial balance to the addresses
+	customAllocs := make([]*genesis.CustomAllocation, 0, len(keys))
+	for _, key := range keys {
+		customAllocs = append(customAllocs, &genesis.CustomAllocation{
+			Address: auth.NewED25519Address(key.PublicKey()),
+			Balance: InitialBalance,
+		})
+	}
+
+	genesis := genesis.NewDefaultGenesis(customAllocs)
+
+	// Set WindowTargetUnits to MaxUint64 for all dimensions to iterate full mempool during block building.
+	genesis.Rules.WindowTargetUnits = fees.Dimensions{math.MaxUint64, math.MaxUint64, math.MaxUint64, math.MaxUint64, math.MaxUint64}
+
+	// Set all limits to MaxUint64 to avoid limiting block size for all dimensions except bandwidth. Must limit bandwidth to avoid building
+	// a block that exceeds the maximum size allowed by AvalancheGo.
+	genesis.Rules.MaxBlockUnits = fees.Dimensions{1800000, math.MaxUint64, math.MaxUint64, math.MaxUint64, math.MaxUint64}
+	genesis.Rules.MinBlockGap = minBlockGap.Milliseconds()
+
+	return genesis
+}
+
+func NewDefaultKeys() []ed25519.PrivateKey {
+	testKeys := make([]ed25519.PrivateKey, len(ed25519HexKeys))
+	for i, keyHex := range ed25519HexKeys {
+		bytes, err := codec.LoadHex(keyHex, ed25519.PrivateKeyLen)
+		if err != nil {
+			panic(err)
+		}
+		testKeys[i] = ed25519.PrivateKey(bytes)
+	}
+
+	return testKeys
+}