From 8e53328335312b5bcd357faecd8558677d5f7540 Mon Sep 17 00:00:00 2001
From: AnkushinDaniil <ankushin.daniil42@gmail.com>
Date: Sun, 22 Dec 2024 17:54:10 +0700
Subject: [PATCH] Add calculateFeeEstimate and TestCalculateFeeEstimate

---
 rpc/simulation.go          | 65 +++++++++++++++++++-------------------
 rpc/simulation_pkg_test.go | 39 +++++++++++++++++++++++
 2 files changed, 71 insertions(+), 33 deletions(-)
 create mode 100644 rpc/simulation_pkg_test.go

diff --git a/rpc/simulation.go b/rpc/simulation.go
index 29e6893180..ec39819e9b 100644
--- a/rpc/simulation.go
+++ b/rpc/simulation.go
@@ -56,7 +56,6 @@ func (h *Handler) SimulateTransactions(id BlockID, transactions []BroadcastedTra
 	return h.simulateTransactions(id, transactions, simulationFlags, false)
 }
 
-//nolint:funlen
 func (h *Handler) simulateTransactions(id BlockID, transactions []BroadcastedTransaction,
 	simulationFlags []SimulationFlag, errOnRevert bool,
 ) ([]SimulatedTransaction, http.Header, *jsonrpc.Error) {
@@ -125,38 +124,7 @@ func (h *Handler) simulateTransactions(id BlockID, transactions []BroadcastedTra
 	for i, overallFee := range overallFees {
 		feeUnit := feeUnit(txns[i])
 
-		var (
-			l1GasPrice     *felt.Felt
-			l2GasPrice     *felt.Felt
-			l1DataGasPrice *felt.Felt
-		)
-
-		switch feeUnit {
-		case FRI:
-			l1GasPrice = header.L1GasPriceSTRK
-			l2GasPrice = header.L2GasPriceSTRK
-			l1DataGasPrice = header.L1DataGasPrice.PriceInFri
-		case WEI:
-			l1GasPrice = header.L1GasPriceETH
-			l2GasPrice = header.L2GasPriceETH
-			l1DataGasPrice = header.L1DataGasPrice.PriceInWei
-		}
-
-		var l1GasConsumed *felt.Felt
-		l1DataGasConsumed := new(felt.Felt).SetUint64(daGas[i].L1DataGas)
-		dataGasFee := new(felt.Felt).Mul(l1DataGasConsumed, l1DataGasPrice)
-		l1GasConsumed = new(felt.Felt).Sub(overallFee, dataGasFee)
-
-		estimate := FeeEstimate{
-			L1GasConsumed:     l1GasConsumed,
-			L2GasConsumed:     &felt.Zero, // TODO: Fix when we have l2 gas price
-			L1GasPrice:        l1GasPrice,
-			L2GasPrice:        l2GasPrice,
-			L1DataGasConsumed: l1DataGasConsumed,
-			L1DataGasPrice:    l1DataGasPrice,
-			OverallFee:        overallFee,
-			Unit:              utils.Ptr(feeUnit),
-		}
+		estimate := calculateFeeEstimate(overallFee, daGas[i].L1DataGas, feeUnit, header)
 
 		trace := traces[i]
 		executionResources := trace.TotalExecutionResources()
@@ -175,6 +143,37 @@ func (h *Handler) simulateTransactions(id BlockID, transactions []BroadcastedTra
 	return result, httpHeader, nil
 }
 
+func calculateFeeEstimate(overallFee *felt.Felt, l1DataGas uint64, feeUnit FeeUnit, header *core.Header) FeeEstimate {
+	var l1GasPrice, l2GasPrice, l1DataGasPrice *felt.Felt
+
+	switch feeUnit {
+	case FRI:
+		l1GasPrice = header.L1GasPriceSTRK
+		l2GasPrice = header.L2GasPriceSTRK
+		l1DataGasPrice = header.L1DataGasPrice.PriceInFri
+	case WEI:
+		l1GasPrice = header.L1GasPriceETH
+		l2GasPrice = header.L2GasPriceETH
+		l1DataGasPrice = header.L1DataGasPrice.PriceInWei
+	}
+
+	l1DataGasConsumed := new(felt.Felt).SetUint64(l1DataGas)
+	dataGasFee := new(felt.Felt).Mul(l1DataGasConsumed, l1DataGasPrice)
+	l1GasConsumed := new(felt.Felt).Sub(overallFee, dataGasFee)
+	l1GasConsumed = l1GasConsumed.Div(l1GasConsumed, l1GasPrice)
+
+	return FeeEstimate{
+		L1GasConsumed:     l1GasConsumed,
+		L2GasConsumed:     &felt.Zero, // TODO: Fix when we have l2 gas price
+		L1GasPrice:        l1GasPrice,
+		L2GasPrice:        l2GasPrice,
+		L1DataGasConsumed: l1DataGasConsumed,
+		L1DataGasPrice:    l1DataGasPrice,
+		OverallFee:        overallFee,
+		Unit:              utils.Ptr(feeUnit),
+	}
+}
+
 type TransactionExecutionErrorData struct {
 	TransactionIndex uint64 `json:"transaction_index"`
 	ExecutionError   string `json:"execution_error"`
diff --git a/rpc/simulation_pkg_test.go b/rpc/simulation_pkg_test.go
new file mode 100644
index 0000000000..44f0c35e37
--- /dev/null
+++ b/rpc/simulation_pkg_test.go
@@ -0,0 +1,39 @@
+package rpc
+
+import (
+	"testing"
+
+	"github.com/NethermindEth/juno/core"
+	"github.com/NethermindEth/juno/core/felt"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestCalculateFeeEstimate(t *testing.T) {
+	l1GasPriceETH := new(felt.Felt).SetUint64(200)
+	l2GasPriceETH := new(felt.Felt).SetUint64(70)
+	l1GasPriceSTRK := new(felt.Felt).SetUint64(100)
+	l2GasPriceSTRK := new(felt.Felt).SetUint64(50)
+	l1DataGasPrice := &core.GasPrice{
+		PriceInWei: new(felt.Felt).SetUint64(10),
+		PriceInFri: new(felt.Felt).SetUint64(5),
+	}
+	header := &core.Header{
+		L1GasPriceETH:  l1GasPriceETH,
+		L2GasPriceETH:  l2GasPriceETH,
+		L1GasPriceSTRK: l1GasPriceSTRK,
+		L2GasPriceSTRK: l2GasPriceSTRK,
+		L1DataGasPrice: l1DataGasPrice,
+	}
+	l1DataGas := uint64(500)
+	overallFee := new(felt.Felt).SetUint64(6000)
+
+	feeEstimate := calculateFeeEstimate(overallFee, l1DataGas, FRI, header)
+
+	assert.Equal(t, l1GasPriceSTRK, feeEstimate.L1GasPrice)
+	assert.Equal(t, l2GasPriceSTRK, feeEstimate.L2GasPrice)
+	assert.Equal(t, l1DataGasPrice.PriceInFri, feeEstimate.L1DataGasPrice)
+	assert.Equal(t, overallFee, feeEstimate.OverallFee)
+	assert.Equal(t, FRI, *feeEstimate.Unit)
+	assert.Equal(t, new(felt.Felt).SetUint64(35), feeEstimate.L1GasConsumed)
+	assert.Equal(t, &felt.Zero, feeEstimate.L2GasConsumed)
+}