fixup: Add panic recovery to load2

Author: maru-ava

tests/load2/generator.go

@@ -14,6 +14,7 @@ import (
 	"golang.org/x/sync/errgroup"
 
 	"github.com/ava-labs/avalanchego/tests"
+	"github.com/ava-labs/avalanchego/utils/logging"
 )
 
 type Test interface {
@@ -61,8 +62,8 @@ func NewLoadGenerator(
 }
 
 func (l LoadGenerator) Run(
-	tc tests.TestContext,
 	ctx context.Context,
+	log logging.Logger,
 	loadTimeout time.Duration,
 	testTimeout time.Duration,
 ) {
@@ -87,10 +88,18 @@ func (l LoadGenerator) Run(
 				ctx, cancel := context.WithTimeout(ctx, testTimeout)
 				defer cancel()
 
-				l.test.Run(tc, ctx, l.wallets[i])
+				execTestWithRecovery(ctx, log, l.test, l.wallets[i])
 			}
 		})
 	}
 
 	_ = eg.Wait()
 }
+
+// execTestWithRecovery ensures assertion-related panics encountered during test execution are recovered
+// and that deferred cleanups are always executed before returning.
+func execTestWithRecovery(ctx context.Context, log logging.Logger, test Test, wallet *Wallet) {
+	tc := tests.NewTestContext(log)
+	defer tc.Recover()
+	test.Run(tc, ctx, wallet)
+}

tests/load2/main/main.go

@@ -50,7 +50,7 @@ func init() {
 func main() {
 	log := tests.NewDefaultLogger("")
 	tc := tests.NewTestContext(log)
-	defer tc.Cleanup()
+	defer tc.RecoverAndExit()
 
 	require := require.New(tc)
 
@@ -116,5 +116,5 @@ func main() {
 	)
 	require.NoError(err)
 
-	generator.Run(tc, ctx, loadTimeout, testTimeout)
+	generator.Run(ctx, log, loadTimeout, testTimeout)
 }