From 08ff8818e4cd4d46e31b681e679ee04cac6a83ac Mon Sep 17 00:00:00 2001
From: skosito <skostic9242@gmail.com>
Date: Tue, 26 Nov 2024 21:08:33 +0000
Subject: [PATCH] test: fix local zetae2e run solana tests and balances (#3217)

* fix local zetae2e run solana tests and balances

* fix config entry for erc20 zrc20

* PR comments
---
 cmd/zetae2e/config/local.yml |  6 +++-
 e2e/runner/balances.go       | 54 ++++++++++++++++++++++++++++++++++--
 2 files changed, 57 insertions(+), 3 deletions(-)

diff --git a/cmd/zetae2e/config/local.yml b/cmd/zetae2e/config/local.yml
index 15de10ceda..88664b4d68 100644
--- a/cmd/zetae2e/config/local.yml
+++ b/cmd/zetae2e/config/local.yml
@@ -3,6 +3,8 @@ default_account:
   bech32_address: "zeta1uhznv7uzyjq84s3q056suc8pkme85lkvhrz3dd"
   evm_address: "0xE5C5367B8224807Ac2207d350E60e1b6F27a7ecC"
   private_key: "d87baf7bf6dc560a252596678c12e41f7d1682837f05b29d411bc3f78ae2c263"
+  solana_address: "37yGiHAnLvWZUNVwu9esp74YQFqxU1qHCbABkDvRddUQ"
+  solana_private_key: "4yqSQxDeTBvn86BuxcN5jmZb2gaobFXrBqu8kiE9rZxNkVMe3LfXmFigRsU4sRp7vk4vVP1ZCFiejDKiXBNWvs2C"
 additional_accounts:
   user_erc20:
     bech32_address: "zeta1datate7xmwm4uk032f9rmcu0cwy7ch7kg6y6zv"
@@ -106,7 +108,7 @@ contracts:
   zevm:
     system_contract: "0x91d18e54DAf4F677cB28167158d6dd21F6aB3921"
     eth_zrc20: "0x13A0c5930C028511Dc02665E7285134B6d11A5f4"
-    erc20_zrc20: "0x48f80608B672DC30DC7e3dbBd0343c5F02C738Eb"
+    erc20_zrc20: "0x0cbe0dF132a6c6B4a2974Fa1b7Fb953CF0Cc798a"
     btc_zrc20: "0xd97B1de3619ed2c6BEb3860147E30cA8A7dC9891"
     uniswap_factory: "0x9fd96203f7b22bCF72d9DCb40ff98302376cE09c"
     uniswap_router: "0x2ca7d64A7EFE2D62A725E2B35Cf7230D6677FfEe"
@@ -115,6 +117,8 @@ contracts:
     test_dapp: "0xA8D5060feb6B456e886F023709A2795373691E63"
     gateway: "0xa825eAa55b497AF892faca73a3797046C10B7c23"
     test_dapp_v2: "0xBFF76e77D56B3C1202107f059425D56f0AEF87Ed"
+    sol_zrc20: "0x48f80608B672DC30DC7e3dbBd0343c5F02C738Eb"
+    spl_zrc20: "0x7c8dDa80bbBE1254a7aACf3219EBe1481c6E01d7"
   evm:
     zeta_eth: "0x733aB8b06DDDEf27Eaa72294B0d7c9cEF7f12db9"
     connector_eth: "0xD28D6A0b8189305551a0A8bd247a6ECa9CE781Ca"
diff --git a/e2e/runner/balances.go b/e2e/runner/balances.go
index 38ac61d70c..8d1fb82597 100644
--- a/e2e/runner/balances.go
+++ b/e2e/runner/balances.go
@@ -7,6 +7,8 @@ import (
 
 	"github.com/btcsuite/btcd/btcutil"
 	"github.com/ethereum/go-ethereum/accounts/abi/bind"
+	"github.com/gagliardetto/solana-go"
+	"github.com/gagliardetto/solana-go/rpc"
 	"github.com/pkg/errors"
 	"github.com/zeta-chain/protocol-contracts/v2/pkg/zrc20.sol"
 )
@@ -25,6 +27,8 @@ type AccountBalances struct {
 	EvmZETA   *big.Int
 	EvmERC20  *big.Int
 	BtcBTC    string
+	SolSOL    *big.Int
+	SolSPL    *big.Int
 }
 
 // AccountBalancesDiff is a struct that contains the difference in the balances of the accounts used in the E2E test
@@ -66,7 +70,7 @@ func (r *E2ERunner) GetAccountBalances(skipBTC bool) (AccountBalances, error) {
 	}
 	zetaSol, err := r.getZRC20BalanceSafe(r.SOLZRC20)
 	if err != nil {
-		r.Logger.Error("get SOL balance: %v", err)
+		return AccountBalances{}, err
 	}
 
 	// evm
@@ -91,6 +95,44 @@ func (r *E2ERunner) GetAccountBalances(skipBTC bool) (AccountBalances, error) {
 		}
 	}
 
+	// solana
+	var solSOL *big.Int
+	var solSPL *big.Int
+	if r.Account.SolanaAddress != "" && r.Account.SolanaPrivateKey != "" {
+		solanaAddr := solana.MustPublicKeyFromBase58(r.Account.SolanaAddress.String())
+		privateKey := solana.MustPrivateKeyFromBase58(r.Account.SolanaPrivateKey.String())
+		solSOLBalance, err := r.SolanaClient.GetBalance(
+			r.Ctx,
+			solanaAddr,
+			rpc.CommitmentFinalized,
+		)
+		if err != nil {
+			return AccountBalances{}, err
+		}
+
+		// #nosec G115 always in range
+		solSOL = big.NewInt(int64(solSOLBalance.Value))
+
+		if r.SPLAddr != (solana.PublicKey{}) {
+			ata := r.ResolveSolanaATA(
+				privateKey,
+				solanaAddr,
+				r.SPLAddr,
+			)
+			splBalance, err := r.SolanaClient.GetTokenAccountBalance(r.Ctx, ata, rpc.CommitmentFinalized)
+			if err != nil {
+				return AccountBalances{}, err
+			}
+
+			solSPLParsed, ok := new(big.Int).SetString(splBalance.Value.Amount, 10)
+			if !ok {
+				return AccountBalances{}, errors.New("can't parse spl balance")
+			}
+
+			solSPL = solSPLParsed
+		}
+	}
+
 	return AccountBalances{
 		ZetaETH:   zetaEth,
 		ZetaZETA:  zetaZeta,
@@ -102,6 +144,8 @@ func (r *E2ERunner) GetAccountBalances(skipBTC bool) (AccountBalances, error) {
 		EvmZETA:   evmZeta,
 		EvmERC20:  evmErc20,
 		BtcBTC:    BtcBTC,
+		SolSOL:    solSOL,
+		SolSPL:    solSPL,
 	}, nil
 }
 
@@ -145,6 +189,7 @@ func (r *E2ERunner) PrintAccountBalances(balances AccountBalances) {
 	r.Logger.Print("* ETH balance:   %s", balances.ZetaETH.String())
 	r.Logger.Print("* ERC20 balance: %s", balances.ZetaERC20.String())
 	r.Logger.Print("* BTC balance:   %s", balances.ZetaBTC.String())
+	r.Logger.Print("* SOL balance: %s", balances.ZetaSOL.String())
 
 	// evm
 	r.Logger.Print("EVM:")
@@ -158,7 +203,12 @@ func (r *E2ERunner) PrintAccountBalances(balances AccountBalances) {
 
 	// solana
 	r.Logger.Print("Solana:")
-	r.Logger.Print("* SOL balance: %s", balances.ZetaSOL.String())
+	if balances.SolSOL != nil {
+		r.Logger.Print("* SOL balance: %s", balances.SolSOL.String())
+	}
+	if balances.SolSPL != nil {
+		r.Logger.Print("* SPL balance: %s", balances.SolSPL.String())
+	}
 }
 
 // PrintTotalDiff shows the difference in the account balances of the accounts used in the e2e test from two balances structs