From 6c5a034587961e5c215fc947952a2d695d3cf8ae Mon Sep 17 00:00:00 2001
From: Maru Newby <maru.newby@avalabs.org>
Date: Fri, 27 Sep 2024 11:40:16 -0700
Subject: [PATCH] [e2e] Add wallet balance logging to staking rewards test

Recent flakes suggest funds exhaustion but reproduction is difficult
without knowing what the balances are.
---
 tests/antithesis/avalanchego/main.go | 18 ++-------------
 tests/e2e/p/staking_rewards.go       |  8 +++++++
 tests/fixture/e2e/helpers.go         | 33 +++++++++++++++++++++++++++-
 3 files changed, 42 insertions(+), 17 deletions(-)

diff --git a/tests/antithesis/avalanchego/main.go b/tests/antithesis/avalanchego/main.go
index 52c19ed3027a..fa3afb7b6e64 100644
--- a/tests/antithesis/avalanchego/main.go
+++ b/tests/antithesis/avalanchego/main.go
@@ -19,6 +19,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/fixture/e2e"
 	"github.com/ava-labs/avalanchego/tests/fixture/tmpnet"
 	"github.com/ava-labs/avalanchego/utils/constants"
 	"github.com/ava-labs/avalanchego/utils/crypto/secp256k1"
@@ -152,22 +153,7 @@ func (w *workload) run(ctx context.Context) {
 	defer tc.Cleanup()
 	require := require.New(tc)
 
-	var (
-		xWallet  = w.wallet.X()
-		xBuilder = xWallet.Builder()
-		pWallet  = w.wallet.P()
-		pBuilder = pWallet.Builder()
-	)
-	xBalances, err := xBuilder.GetFTBalance()
-	require.NoError(err, "failed to fetch X-chain balances")
-	pBalances, err := pBuilder.GetBalance()
-	require.NoError(err, "failed to fetch P-chain balances")
-	var (
-		xContext    = xBuilder.Context()
-		avaxAssetID = xContext.AVAXAssetID
-		xAVAX       = xBalances[avaxAssetID]
-		pAVAX       = pBalances[avaxAssetID]
-	)
+	xAVAX, pAVAX := e2e.GetWalletBalances(tc, w.wallet)
 	log.Printf("wallet starting with %d X-chain nAVAX and %d P-chain nAVAX", xAVAX, pAVAX)
 	assert.Reachable("wallet starting", map[string]any{
 		"worker":   w.id,
diff --git a/tests/e2e/p/staking_rewards.go b/tests/e2e/p/staking_rewards.go
index 84156a9b74d4..29545450ddbc 100644
--- a/tests/e2e/p/staking_rewards.go
+++ b/tests/e2e/p/staking_rewards.go
@@ -111,6 +111,8 @@ var _ = ginkgo.Describe("[Staking Rewards]", func() {
 		require.NoError(err)
 
 		tc.By("adding alpha node as a validator", func() {
+			e2e.OutputWalletBalances(tc, baseWallet)
+
 			endTime := time.Now().Add(targetValidationPeriod)
 			tc.Outf("validation period ending at: %v\n", endTime)
 
@@ -143,6 +145,8 @@ var _ = ginkgo.Describe("[Staking Rewards]", func() {
 		tc.Outf("beta node validation period ending at: %v\n", betaValidatorEndTime)
 
 		tc.By("adding beta node as a validator", func() {
+			e2e.OutputWalletBalances(tc, baseWallet)
+
 			_, err := pWallet.IssueAddPermissionlessValidatorTx(
 				&txs.SubnetValidator{
 					Validator: txs.Validator{
@@ -173,6 +177,8 @@ var _ = ginkgo.Describe("[Staking Rewards]", func() {
 		require.NoError(err)
 
 		tc.By("adding gamma as a delegator to the alpha node", func() {
+			e2e.OutputWalletBalances(tc, baseWallet)
+
 			endTime := time.Now().Add(targetDelegationPeriod)
 			tc.Outf("delegation period ending at: %v\n", endTime)
 
@@ -196,6 +202,8 @@ var _ = ginkgo.Describe("[Staking Rewards]", func() {
 		})
 
 		tc.By("adding delta as delegator to the beta node", func() {
+			e2e.OutputWalletBalances(tc, baseWallet)
+
 			endTime := time.Now().Add(targetDelegationPeriod)
 			tc.Outf("delegation period ending at: %v\n", endTime)
 
diff --git a/tests/fixture/e2e/helpers.go b/tests/fixture/e2e/helpers.go
index f26795c52801..ba06c60e3a38 100644
--- a/tests/fixture/e2e/helpers.go
+++ b/tests/fixture/e2e/helpers.go
@@ -68,7 +68,7 @@ func NewWallet(tc tests.TestContext, keychain *secp256k1fx.Keychain, nodeURI tmp
 		EthKeychain:  keychain,
 	})
 	require.NoError(tc, err)
-	return primary.NewWalletWithOptions(
+	wallet := primary.NewWalletWithOptions(
 		baseWallet,
 		common.WithPostIssuanceFunc(
 			func(id ids.ID) {
@@ -76,6 +76,37 @@ func NewWallet(tc tests.TestContext, keychain *secp256k1fx.Keychain, nodeURI tmp
 			},
 		),
 	)
+	xAVAX, pAVAX := GetWalletBalances(tc, wallet)
+	tc.Outf("{{blue}}  wallet starting with %d X-chain nAVAX and %d P-chain nAVAX{{/}}\n", xAVAX, pAVAX)
+	return wallet
+}
+
+// OutputWalletBalances outputs the X-Chain and P-Chain balances of the provided wallet.
+func OutputWalletBalances(tc tests.TestContext, wallet primary.Wallet) {
+	xAVAX, pAVAX := GetWalletBalances(tc, wallet)
+	tc.Outf("{{blue}}  wallet has %d X-chain nAVAX and %d P-chain nAVAX{{/}}\n", xAVAX, pAVAX)
+}
+
+// GetWalletBalances retrieves the X-Chain and P-Chain balances of the provided wallet.
+func GetWalletBalances(tc tests.TestContext, wallet primary.Wallet) (uint64, uint64) {
+	require := require.New(tc)
+	var (
+		xWallet  = wallet.X()
+		xBuilder = xWallet.Builder()
+		pWallet  = wallet.P()
+		pBuilder = pWallet.Builder()
+	)
+	xBalances, err := xBuilder.GetFTBalance()
+	require.NoError(err, "failed to fetch X-chain balances")
+	pBalances, err := pBuilder.GetBalance()
+	require.NoError(err, "failed to fetch P-chain balances")
+	var (
+		xContext    = xBuilder.Context()
+		avaxAssetID = xContext.AVAXAssetID
+		xAVAX       = xBalances[avaxAssetID]
+		pAVAX       = pBalances[avaxAssetID]
+	)
+	return xAVAX, pAVAX
 }
 
 // Create a new eth client targeting the specified node URI.