From 103ad5b8df9dc2e46b4cea8a456c4b47d7872adc Mon Sep 17 00:00:00 2001
From: Gabriel de Quadros Ligneul <gligneul@offchainlabs.com>
Date: Thu, 3 Oct 2024 11:38:46 -0300
Subject: [PATCH] Replace util function with existing TestClient one

---
 system_tests/common_test.go      | 13 ++++++++-----
 system_tests/program_gas_test.go | 26 ++++++++++----------------
 2 files changed, 18 insertions(+), 21 deletions(-)

diff --git a/system_tests/common_test.go b/system_tests/common_test.go
index 67bc01c92d..93c38b5eae 100644
--- a/system_tests/common_test.go
+++ b/system_tests/common_test.go
@@ -138,8 +138,8 @@ func (tc *TestClient) GetBaseFeeAt(t *testing.T, blockNum *big.Int) *big.Int {
 	return GetBaseFeeAt(t, tc.Client, tc.ctx, blockNum)
 }
 
-func (tc *TestClient) SendWaitTestTransactions(t *testing.T, txs []*types.Transaction) {
-	SendWaitTestTransactions(t, tc.ctx, tc.Client, txs)
+func (tc *TestClient) SendWaitTestTransactions(t *testing.T, txs []*types.Transaction) []*types.Receipt {
+	return SendWaitTestTransactions(t, tc.ctx, tc.Client, txs)
 }
 
 func (tc *TestClient) DeploySimple(t *testing.T, auth bind.TransactOpts) (common.Address, *mocksgen.Simple) {
@@ -763,15 +763,18 @@ func (b *NodeBuilder) BridgeBalance(t *testing.T, account string, amount *big.In
 	return BridgeBalance(t, account, amount, b.L1Info, b.L2Info, b.L1.Client, b.L2.Client, b.ctx)
 }
 
-func SendWaitTestTransactions(t *testing.T, ctx context.Context, client *ethclient.Client, txs []*types.Transaction) {
+func SendWaitTestTransactions(t *testing.T, ctx context.Context, client *ethclient.Client, txs []*types.Transaction) []*types.Receipt {
 	t.Helper()
+	receipts := make([]*types.Receipt, len(txs))
 	for _, tx := range txs {
 		Require(t, client.SendTransaction(ctx, tx))
 	}
-	for _, tx := range txs {
-		_, err := EnsureTxSucceeded(ctx, client, tx)
+	for i, tx := range txs {
+		var err error
+		receipts[i], err = EnsureTxSucceeded(ctx, client, tx)
 		Require(t, err)
 	}
+	return receipts
 }
 
 func TransferBalance(
diff --git a/system_tests/program_gas_test.go b/system_tests/program_gas_test.go
index 063a21d3b0..fdcd53eeac 100644
--- a/system_tests/program_gas_test.go
+++ b/system_tests/program_gas_test.go
@@ -291,14 +291,18 @@ func compareGasUsage(
 	}
 
 	const txGas uint64 = 32_000_000
-	tx := builder.L2Info.PrepareTxTo("Owner", &evmContract, txGas, txValue, txData)
-	evmGas := sendAndEnsureTransaction(t, builder.ctx, builder.L2.Client, tx)
-	evmGasUsage, err := evmOpcodesGasUsage(builder.ctx, builder.L2.Client.Client(), tx)
+	txs := []*types.Transaction{
+		builder.L2Info.PrepareTxTo("Owner", &evmContract, txGas, txValue, txData),
+		builder.L2Info.PrepareTxTo("Owner", &stylusContract, txGas, txValue, txData),
+	}
+	receipts := builder.L2.SendWaitTestTransactions(t, txs)
+
+	evmGas := receipts[0].GasUsedForL2()
+	evmGasUsage, err := evmOpcodesGasUsage(builder.ctx, builder.L2.Client.Client(), txs[0])
 	Require(t, err)
 
-	tx = builder.L2Info.PrepareTxTo("Owner", &stylusContract, txGas, txValue, txData)
-	stylusGas := sendAndEnsureTransaction(t, builder.ctx, builder.L2.Client, tx)
-	stylusGasUsage, err := stylusHostiosGasUsage(builder.ctx, builder.L2.Client.Client(), tx)
+	stylusGas := receipts[1].GasUsedForL2()
+	stylusGasUsage, err := stylusHostiosGasUsage(builder.ctx, builder.L2.Client.Client(), txs[1])
 	Require(t, err)
 
 	t.Logf("evm total usage: %v - stylus total usage: %v", evmGas, stylusGas)
@@ -415,13 +419,3 @@ func checkPercentDiff(t *testing.T, a, b float64, maxAllowedDifference float64)
 		Fatal(t, fmt.Sprintf("gas usages are too different; got %v, max allowed is %v", percentageDifference, maxAllowedDifference))
 	}
 }
-
-// sendAndEnsureTransaction sends a transaction, ensures it succeed, and returns the total gas cost.
-func sendAndEnsureTransaction(t *testing.T, ctx context.Context, client *ethclient.Client, tx *types.Transaction) uint64 {
-	t.Helper()
-	err := client.SendTransaction(ctx, tx)
-	Require(t, err)
-	receipt, err := EnsureTxSucceeded(ctx, client, tx)
-	Require(t, err)
-	return receipt.GasUsedForL2()
-}