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

tests/antithesis/avalanchego/main.go

@@ -6,6 +6,7 @@ package main
 import (
 	"context"
 	"crypto/rand"
+	"fmt"
 	"log"
 	"math/big"
 	"time"
@@ -19,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"
@@ -42,9 +44,12 @@ 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() {
-	tc := tests.NewTestContext()
-	defer tc.Cleanup()
+	tc := antithesis.NewInstrumentedTestContext()
+	defer tc.RecoverAndExit()
 	require := require.New(tc)
 
 	c := antithesis.NewConfig(
@@ -54,16 +59,13 @@ 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, &primary.WalletConfig{
-		URI:          c.URIs[0],
-		AVAXKeychain: kc,
-		EthKeychain:  kc,
-	})
-	require.NoError(err, "failed to initialize wallet")
-	log.Printf("synced wallet in %s", time.Since(walletSyncStartTime))
+	wallet := e2e.NewWallet(tc, kc, tmpnet.NodeURI{URI: c.URIs[0]})
+	tc.Outf("synced wallet in %s", time.Since(walletSyncStartTime))
 
 	genesisWorkload := &workload{
 		id:     0,
@@ -104,21 +106,15 @@ func main() {
 			},
 		}})
 		require.NoError(err, "failed to issue initial funding X-chain baseTx")
-		log.Printf("issued initial funding X-chain baseTx %s in %s", baseTx.ID(), time.Since(baseStartTime))
+		tc.Outf("issued initial funding X-chain baseTx %s in %s", baseTx.ID(), time.Since(baseStartTime))
 
-		// TODO(marun) Enable cleanup of these contexts
 		genesisWorkload.confirmXChainTx(ctx, baseTx)
 
 		uri := c.URIs[i%len(c.URIs)]
 		kc := secp256k1fx.NewKeychain(key)
 		walletSyncStartTime := time.Now()
-		wallet, err := primary.MakeWallet(ctx, &primary.WalletConfig{
-			URI:          uri,
-			AVAXKeychain: kc,
-			EthKeychain:  kc,
-		})
-		require.NoError(err, "failed to initialize wallet")
-		log.Printf("synced wallet in %s", time.Since(walletSyncStartTime))
+		wallet := e2e.NewWallet(tc, kc, tmpnet.NodeURI{URI: uri})
+		tc.Outf("synced wallet in %s", time.Since(walletSyncStartTime))
 
 		workloads[i] = &workload{
 			id:     i,
@@ -146,44 +142,54 @@ 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,
+		fmt.Sprintf("worker %d: ", w.id),
+		map[string]any{
+			"worker": w.id,
+		},
+	)
+}
+
 func (w *workload) run(ctx context.Context) {
 	timer := timerpkg.StoppedTimer()
 
-	tc := tests.NewTestContext()
-	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)
-	log.Printf("wallet starting with %d X-chain nAVAX and %d P-chain nAVAX", xAVAX, pAVAX)
+	tc.Outf("wallet starting with %d X-chain nAVAX and %d P-chain nAVAX", w.id, xAVAX, pAVAX)
 	assert.Reachable("wallet starting", map[string]any{
 		"worker":   w.id,
 		"xBalance": xAVAX,
 		"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()
-		log.Printf("wallet %d executing flow %d", w.id, 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 {
+			tc.Outf("failed to read randomness: %s", err)
+			assert.Unreachable("failed to read randomness", map[string]any{
+				"worker": w.id,
+				"err":    err,
+			})
+			rawExecutionDelay = defaultExecutionDelay
+		}
+		executionDelay := time.Duration(rawExecutionDelay.Int64())
+		tc.Outf("waiting %v", executionDelay)
+		timer.Reset(executionDelay)
 
-		timer.Reset(time.Duration(val.Int64()))
 		select {
 		case <-ctx.Done():
 			return
@@ -192,6 +198,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}`
+		tc.Outf("executing issueXChainBaseTx")
+		w.issueXChainBaseTx(ctx)
+	case 1:
+		tc.Outf("executing issueXChainCreateAssetTx")
+		w.issueXChainCreateAssetTx(ctx)
+	case 2:
+		tc.Outf("executing issueXChainOperationTx")
+		w.issueXChainOperationTx(ctx)
+	case 3:
+		tc.Outf("executing issueXToPTransfer")
+		w.issueXToPTransfer(ctx)
+	case 4:
+		tc.Outf("executing issuePToXTransfer")
+		w.issuePToXTransfer(ctx)
+	case 5:
+		tc.Outf("executing banff.TestCustomAssetTransfer")
+		addr, _ := w.addrs.Peek()
+		banff.TestCustomAssetTransfer(tc, w.wallet, addr)
+	case 6:
+		tc.Outf("sleeping")
+	}
+}
+
 func (w *workload) issueXChainBaseTx(ctx context.Context) {
 	var (
 		xWallet  = w.wallet.X()

tests/antithesis/context.go

@@ -0,0 +1,39 @@
+// 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"
+)
+
+// NewInstrumentedTestContext returns a test context that makes antithesis SDK assertions.
+func NewInstrumentedTestContext() *tests.SimpleTestContext {
+	return NewInstrumentedTestContextWithArgs(nil, "", nil)
+}
+
+// NewInstrumentedTestContextWithArgs returns a test context that makes antithesis SDK assertions.
+func NewInstrumentedTestContextWithArgs(
+	ctx context.Context,
+	prefix string,
+	details map[string]any,
+) *tests.SimpleTestContext {
+	return tests.NewTestContextWithArgs(
+		ctx,
+		prefix,
+		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

@@ -6,6 +6,7 @@ package main
 import (
 	"context"
 	"crypto/rand"
+	"fmt"
 	"log"
 	"math/big"
 	"time"
@@ -36,8 +37,8 @@ const (
 )
 
 func main() {
-	tc := tests.NewTestContext()
-	defer tc.Cleanup()
+	tc := antithesis.NewInstrumentedTestContext()
+	defer tc.RecoverAndExit()
 	require := require.New(tc)
 
 	c := antithesis.NewConfigWithSubnets(
@@ -52,6 +53,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)
 	log.Printf("CHAIN IDS: %v", c.ChainIDs)
@@ -127,8 +130,16 @@ type workload struct {
 func (w *workload) run(ctx context.Context) {
 	timer := timerpkg.StoppedTimer()
 
-	tc := tests.NewTestContext()
-	defer tc.Cleanup()
+	tc := antithesis.NewInstrumentedTestContextWithArgs(
+		ctx,
+		fmt.Sprintf("worker %d: ", w.id),
+		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
 }