From 1a2b437eb5598830cfa68ee403ffc8dd8ef22f78 Mon Sep 17 00:00:00 2001
From: skosito <skostic9242@gmail.com>
Date: Tue, 26 Nov 2024 13:32:03 +0100
Subject: [PATCH] fix local zetae2e run solana tests and balances

---
 cmd/zetae2e/config/local.yml |  4 +++
 e2e/runner/balances.go       | 47 ++++++++++++++++++++++++++++++++++--
 2 files changed, 49 insertions(+), 2 deletions(-)

diff --git a/cmd/zetae2e/config/local.yml b/cmd/zetae2e/config/local.yml
index 15de10ceda..75f6a0515c 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"
@@ -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..de57af118c 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,41 @@ func (r *E2ERunner) GetAccountBalances(skipBTC bool) (AccountBalances, error) {
 		}
 	}
 
+	// solana
+	var solSOL *big.Int
+	var solSPL *big.Int
+	if r.Account.SolanaAddress != "" && r.Account.SolanaPrivateKey != "" {
+		solSOLBalance, err := r.SolanaClient.GetBalance(
+			r.Ctx,
+			solana.MustPublicKeyFromBase58(r.Account.SolanaAddress.String()),
+			rpc.CommitmentFinalized,
+		)
+		if err != nil {
+			return AccountBalances{}, err
+		}
+
+		solSOL = big.NewInt(int64(solSOLBalance.Value))
+
+		if r.SPLAddr != (solana.PublicKey{}) {
+			ata := r.ResolveSolanaATA(
+				solana.MustPrivateKeyFromBase58(r.Account.SolanaPrivateKey.String()),
+				solana.MustPublicKeyFromBase58(r.Account.SolanaAddress.String()),
+				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 +141,8 @@ func (r *E2ERunner) GetAccountBalances(skipBTC bool) (AccountBalances, error) {
 		EvmZETA:   evmZeta,
 		EvmERC20:  evmErc20,
 		BtcBTC:    BtcBTC,
+		SolSOL:    solSOL,
+		SolSPL:    solSPL,
 	}, nil
 }
 
@@ -145,6 +186,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 +200,8 @@ func (r *E2ERunner) PrintAccountBalances(balances AccountBalances) {
 
 	// solana
 	r.Logger.Print("Solana:")
-	r.Logger.Print("* SOL balance: %s", balances.ZetaSOL.String())
+	r.Logger.Print("* SOL balance: %s", balances.SolSOL.String())
+	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