Implement Wait For Transaciton Receipt in accounts

NethermindEth · Oct 4, 2023 · c16b4ca · c16b4ca
1 parent a381df5
commit c16b4ca
Show file tree

Hide file tree

Showing 3 changed files with 70 additions and 22 deletions.
diff --git a/account/account.go b/account/account.go
@@ -3,6 +3,7 @@ package account
 import (
 	"context"
 	"errors"
+	"time"
 
 	"github.com/NethermindEth/juno/core/felt"
 	starknetgo "github.com/NethermindEth/starknet.go"
@@ -35,6 +36,7 @@ type AccountInterface interface {
 	SignDeployAccountTransaction(ctx context.Context, tx *rpc.DeployAccountTxn, precomputeAddress *felt.Felt) error
 	SignDeclareTransaction(ctx context.Context, tx *rpc.DeclareTxnV2) error
 	PrecomputeAddress(deployerAddress *felt.Felt, salt *felt.Felt, classHash *felt.Felt, constructorCalldata []*felt.Felt) (*felt.Felt, error)
+	WaitForTransactionReceipt(ctx context.Context, transactionHash *felt.Felt, pollInterval time.Duration) (*rpc.TransactionReceipt, error)
 }
 
 var _ AccountInterface = &Account{}
@@ -293,6 +295,27 @@ func (account *Account) PrecomputeAddress(deployerAddress *felt.Felt, salt *felt
 
 }
 
+// WaitForTransactionReceipt waits for the transaction to succeed or fail
+func (account *Account) WaitForTransactionReceipt(ctx context.Context, transactionHash *felt.Felt, pollInterval time.Duration) (*rpc.TransactionReceipt, error) {
+	t := time.NewTicker(pollInterval)
+	for {
+		select {
+		case <-ctx.Done():
+			return nil, ctx.Err()
+		case <-t.C:
+			receipt, err := account.TransactionReceipt(ctx, transactionHash)
+			if err != nil {
+				if err.Error() == rpc.ErrHashNotFound.Error() {
+					continue
+				} else {
+					return nil, err
+				}
+			}
+			return &receipt, nil
+		}
+	}
+}
+
 // BuildInvokeTx formats the calldata and signs the transaction
 func (account *Account) BuildInvokeTx(ctx context.Context, invokeTx *rpc.InvokeTxnV1, fnCall *[]rpc.FunctionCall) error {
 	invokeTx.Calldata = FmtCalldata(*fnCall)

diff --git a/account/account_test.go b/account/account_test.go
@@ -2,11 +2,13 @@ package account_test
 
 import (
 	"context"
+	"errors"
 	"flag"
 	"fmt"
 	"math/big"
 	"os"
 	"testing"
+	"time"
 
 	"github.com/NethermindEth/juno/core/felt"
 	starknetgo "github.com/NethermindEth/starknet.go"
@@ -530,6 +532,51 @@ func TestTransactionHashDeclare(t *testing.T) {
 	require.Equal(t, expectedHash.String(), hash.String(), "TransactionHashDeclare not what expected")
 }
 
+func TestWaitForTransactionReceiptMOCK(t *testing.T) {
+	mockCtrl := gomock.NewController(t)
+	t.Cleanup(mockCtrl.Finish)
+	mockRpcProvider := mocks.NewMockRpcProvider(mockCtrl)
+
+	mockRpcProvider.EXPECT().ChainID(context.Background()).Return("SN_GOERLI", nil)
+	acnt, err := account.NewAccount(mockRpcProvider, &felt.Zero, "", starknetgo.NewMemKeystore())
+	require.NoError(t, err, "error returned from account.NewAccount()")
+
+	type testSetType struct {
+		Hash            *felt.Felt
+		ExpectedErr     error
+		ExpectedReceipt rpc.TransactionReceipt
+	}
+	testSet := map[string][]testSetType{
+		"mock": {
+			{
+				Hash:            new(felt.Felt).SetUint64(1),
+				ExpectedReceipt: nil,
+				ExpectedErr:     errors.New("UnExpectedErr"),
+			},
+			{
+				Hash: new(felt.Felt).SetUint64(2),
+				ExpectedReceipt: rpc.InvokeTransactionReceipt{
+					TransactionHash: new(felt.Felt).SetUint64(2),
+					ExecutionStatus: rpc.TxnExecutionStatusSUCCEEDED,
+				},
+				ExpectedErr: nil,
+			},
+		},
+	}[testEnv]
+
+	for _, test := range testSet {
+		mockRpcProvider.EXPECT().TransactionReceipt(context.Background(), test.Hash).Return(test.ExpectedReceipt, test.ExpectedErr)
+		resp, err := acnt.WaitForTransactionReceipt(context.Background(), test.Hash, 2*time.Second)
+
+		if err != nil {
+			require.Equal(t, test.ExpectedErr, err)
+		} else {
+			require.Equal(t, test.ExpectedReceipt.GetExecutionStatus(), (*resp).GetExecutionStatus())
+		}
+	}
+
+}
+
 func newDevnet(t *testing.T, url string) ([]test.TestAccount, error) {
 	devnet := test.NewDevNet(url)
 	acnts, err := devnet.Accounts()

diff --git a/rpc/transaction.go b/rpc/transaction.go
@@ -5,7 +5,6 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
-	"time"
 
 	"github.com/NethermindEth/juno/core/felt"
 )
@@ -107,24 +106,3 @@ func (provider *Provider) TransactionReceipt(ctx context.Context, transactionHas
 	}
 	return receipt.TransactionReceipt, nil
 }
-
-// WaitForTransaction waits for the transaction to succeed or fail
-func (provider *Provider) WaitForTransaction(ctx context.Context, transactionHash *felt.Felt, pollInterval time.Duration) (TxnExecutionStatus, error) {
-	t := time.NewTicker(pollInterval)
-	for {
-		select {
-		case <-ctx.Done():
-			return "", ctx.Err()
-		case <-t.C:
-			_, err := provider.TransactionByHash(ctx, transactionHash)
-			if err != nil {
-				break
-			}
-			receipt, err := provider.TransactionReceipt(ctx, transactionHash)
-			if err != nil {
-				continue
-			}
-			return receipt.GetExecutionStatus(), nil
-		}
-	}
-}