[antithesis] Enable reuse of banff e2e test for antithesis testing

Author: maru-ava

tests/antithesis/avalanchego/main.go

@@ -20,6 +20,7 @@ import (
 	"github.com/ava-labs/avalanchego/ids"
 	"github.com/ava-labs/avalanchego/tests"
 	"github.com/ava-labs/avalanchego/tests/antithesis"
+	"github.com/ava-labs/avalanchego/tests/e2e/banff"
 	"github.com/ava-labs/avalanchego/tests/fixture/e2e"
 	"github.com/ava-labs/avalanchego/tests/fixture/tmpnet"
 	"github.com/ava-labs/avalanchego/utils/constants"
@@ -44,10 +45,13 @@ import (
 
 const NumKeys = 5
 
+// TODO(marun) Switch to using zap for logging
+// TODO(marun) Extract the common elements of test execution for reuse across test setups
+
 func main() {
 	// TODO(marun) Support choosing the log format
-	tc := tests.NewTestContext(tests.NewDefaultLogger(""))
-	defer tc.Cleanup()
+	tc := antithesis.NewInstrumentedTestContext(tests.NewDefaultLogger(""))
+	defer tc.RecoverAndExit()
 	require := require.New(tc)
 
 	c := antithesis.NewConfig(
@@ -57,17 +61,12 @@ func main() {
 		},
 	)
 	ctx := tests.DefaultNotifyContext(c.Duration, tc.DeferCleanup)
+	// Ensure contexts sourced from the test context use the notify context as their parent
+	tc.SetDefaultContextParent(ctx)
 
 	kc := secp256k1fx.NewKeychain(genesis.EWOQKey)
 	walletSyncStartTime := time.Now()
-	wallet, err := primary.MakeWallet(
-		ctx,
-		c.URIs[0],
-		kc,
-		kc,
-		primary.WalletConfig{},
-	)
-	require.NoError(err, "failed to initialize wallet")
+	wallet := e2e.NewWallet(tc, kc, tmpnet.NodeURI{URI: c.URIs[0]})
 	tc.Log().Info("synced wallet",
 		zap.Duration("duration", time.Since(walletSyncStartTime)),
 	)
@@ -122,14 +121,7 @@ func main() {
 		uri := c.URIs[i%len(c.URIs)]
 		kc := secp256k1fx.NewKeychain(key)
 		walletSyncStartTime := time.Now()
-		wallet, err := primary.MakeWallet(
-			ctx,
-			uri,
-			kc,
-			kc,
-			primary.WalletConfig{},
-		)
-		require.NoError(err, "failed to initialize wallet")
+		wallet := e2e.NewWallet(tc, kc, tmpnet.NodeURI{URI: uri})
 		tc.Log().Info("synced wallet",
 			zap.Duration("duration", time.Since(walletSyncStartTime)),
 		)
@@ -162,12 +154,25 @@ type workload struct {
 	uris   []string
 }
 
+// newTestContext returns a test context that ensures that log output and assertions are
+// associated with this worker.
+func (w *workload) newTestContext(ctx context.Context) *tests.SimpleTestContext {
+	return antithesis.NewInstrumentedTestContextWithArgs(
+		ctx,
+		w.log,
+		map[string]any{
+			"worker": w.id,
+		},
+	)
+}
+
 func (w *workload) run(ctx context.Context) {
 	timer := timerpkg.StoppedTimer()
 
-	tc := tests.NewTestContext(w.log)
-	defer tc.Cleanup()
-	require := require.New(tc)
+	tc := w.newTestContext(ctx)
+	// Any assertion failure from this test context will result in process exit due to the
+	// panic being rethrown. This ensures that failures in test setup are fatal.
+	defer tc.Recover(true /* rethrow */)
 
 	xAVAX, pAVAX := e2e.GetWalletBalances(tc, w.wallet)
 	assert.Reachable("wallet starting", map[string]any{
@@ -176,32 +181,30 @@ func (w *workload) run(ctx context.Context) {
 		"pBalance": pAVAX,
 	})
 
+	defaultExecutionDelay := big.NewInt(int64(time.Second))
 	for {
-		val, err := rand.Int(rand.Reader, big.NewInt(5))
-		require.NoError(err, "failed to read randomness")
-
-		flowID := val.Int64()
-		w.log.Info("executing test",
-			zap.Int("workerID", w.id),
-			zap.Int64("flowID", flowID),
-		)
-		switch flowID {
-		case 0:
-			w.issueXChainBaseTx(ctx)
-		case 1:
-			w.issueXChainCreateAssetTx(ctx)
-		case 2:
-			w.issueXChainOperationTx(ctx)
-		case 3:
-			w.issueXToPTransfer(ctx)
-		case 4:
-			w.issuePToXTransfer(ctx)
-		}
+		w.executeTest(ctx)
 
-		val, err = rand.Int(rand.Reader, big.NewInt(int64(time.Second)))
-		require.NoError(err, "failed to read randomness")
+		// Delay execution of the next test by a random duration
+		rawExecutionDelay, err := rand.Int(rand.Reader, defaultExecutionDelay)
+		// Avoid using require.NoError since the execution delay is not critical and an
+		// assertion failure in this function is fatal.
+		if err != nil {
+			w.log.Error("failed to read randomness",
+				zap.Error(err),
+			)
+			assert.Unreachable("failed to read randomness", map[string]any{
+				"worker": w.id,
+				"err":    err,
+			})
+			rawExecutionDelay = defaultExecutionDelay
+		}
+		executionDelay := time.Duration(rawExecutionDelay.Int64())
+		w.log.Info("waiting",
+			zap.Duration("duration", executionDelay),
+		)
+		timer.Reset(executionDelay)
 
-		timer.Reset(time.Duration(val.Int64()))
 		select {
 		case <-ctx.Done():
 			return
@@ -210,6 +213,44 @@ func (w *workload) run(ctx context.Context) {
 	}
 }
 
+// executeTest executes a test at random.
+func (w *workload) executeTest(ctx context.Context) {
+	tc := w.newTestContext(ctx)
+	// Panics will be recovered without being rethrown, ensuring that test failures are not fatal.
+	defer tc.Recover(false /* rethrow */)
+	require := require.New(tc)
+
+	val, err := rand.Int(rand.Reader, big.NewInt(6))
+	require.NoError(err, "failed to read randomness")
+
+	// TODO(marun)
+	flowID := val.Int64()
+	switch flowID {
+	case 0:
+		// TODO(marun) Create abstraction for a test that supports a name e.g. `aTest{name: "foo", mytestfunc}`
+		w.log.Info("executing issueXChainBaseTx")
+		w.issueXChainBaseTx(ctx)
+	case 1:
+		w.log.Info("executing issueXChainCreateAssetTx")
+		w.issueXChainCreateAssetTx(ctx)
+	case 2:
+		w.log.Info("executing issueXChainOperationTx")
+		w.issueXChainOperationTx(ctx)
+	case 3:
+		w.log.Info("executing issueXToPTransfer")
+		w.issueXToPTransfer(ctx)
+	case 4:
+		w.log.Info("executing issuePToXTransfer")
+		w.issuePToXTransfer(ctx)
+	case 5:
+		w.log.Info("executing banff.TestCustomAssetTransfer")
+		addr, _ := w.addrs.Peek()
+		banff.TestCustomAssetTransfer(tc, w.wallet, addr)
+	case 6:
+		w.log.Info("sleeping")
+	}
+}
+
 func (w *workload) issueXChainBaseTx(ctx context.Context) {
 	var (
 		xWallet  = w.wallet.X()

tests/antithesis/context.go

@@ -0,0 +1,40 @@
+// Copyright (C) 2019-2024, Ava Labs, Inc. All rights reserved.
+// See the file LICENSE for licensing terms.
+
+package antithesis
+
+import (
+	"context"
+	"fmt"
+	"maps"
+
+	"github.com/antithesishq/antithesis-sdk-go/assert"
+
+	"github.com/ava-labs/avalanchego/tests"
+	"github.com/ava-labs/avalanchego/utils/logging"
+)
+
+// NewInstrumentedTestContext returns a test context that makes antithesis SDK assertions.
+func NewInstrumentedTestContext(log logging.Logger) *tests.SimpleTestContext {
+	return NewInstrumentedTestContextWithArgs(context.Background(), log, nil)
+}
+
+// NewInstrumentedTestContextWithArgs returns a test context that makes antithesis SDK assertions.
+func NewInstrumentedTestContextWithArgs(
+	ctx context.Context,
+	log logging.Logger,
+	details map[string]any,
+) *tests.SimpleTestContext {
+	return tests.NewTestContextWithArgs(
+		ctx,
+		log,
+		func(format string, args ...any) {
+			assert.Unreachable(fmt.Sprintf("Assertion failure: "+format, args...), details)
+		},
+		func(r any) {
+			detailsClone := maps.Clone(details)
+			details["panic"] = r
+			assert.Unreachable("unexpected panic", detailsClone)
+		},
+	)
+}

tests/antithesis/xsvm/main.go

@@ -39,8 +39,8 @@ const (
 
 func main() {
 	// TODO(marun) Support choosing the log format
-	tc := tests.NewTestContext(tests.NewDefaultLogger(""))
-	defer tc.Cleanup()
+	tc := antithesis.NewInstrumentedTestContext(tests.NewDefaultLogger(""))
+	defer tc.RecoverAndExit()
 	require := require.New(tc)
 
 	c := antithesis.NewConfigWithSubnets(
@@ -55,6 +55,8 @@ func main() {
 		},
 	)
 	ctx := tests.DefaultNotifyContext(c.Duration, tc.DeferCleanup)
+	// Ensure contexts sourced from the test context use the notify context as their parent
+	tc.SetDefaultContextParent(ctx)
 
 	require.Len(c.ChainIDs, 1)
 	tc.Log().Debug("raw chain ID",
@@ -140,8 +142,16 @@ type workload struct {
 func (w *workload) run(ctx context.Context) {
 	timer := timerpkg.StoppedTimer()
 
-	tc := tests.NewTestContext(w.log)
-	defer tc.Cleanup()
+	tc := antithesis.NewInstrumentedTestContextWithArgs(
+		ctx,
+		w.log,
+		map[string]any{
+			"worker": w.id,
+		},
+	)
+	// Any assertion failure from this test context will result in process exit due to the
+	// panic being rethrown. This ensures that failures in test setup are fatal.
+	defer tc.Recover(true /* rethrow */)
 	require := require.New(tc)
 
 	uri := w.uris[w.id%len(w.uris)]

tests/context_helpers.go

@@ -16,7 +16,11 @@ const DefaultTimeout = 2 * time.Minute
 
 // Helper simplifying use of a timed context by canceling the context with the test context.
 func ContextWithTimeout(tc TestContext, duration time.Duration) context.Context {
-	ctx, cancel := context.WithTimeout(context.Background(), duration)
+	parent := tc.GetDefaultContextParent()
+	if parent == nil {
+		parent = context.Background()
+	}
+	ctx, cancel := context.WithTimeout(parent, duration)
 	tc.DeferCleanup(cancel)
 	return ctx
 }