From 1b4b77b02b6728ec1702e16ae2858051d9866620 Mon Sep 17 00:00:00 2001
From: swarmHaseeb <haseeb@ethswarm.org>
Date: Fri, 16 Dec 2022 21:09:52 +0500
Subject: [PATCH] refactor: minor changes and bug fixes

---
 pkg/api/export_test.go                    |  1 +
 pkg/api/staking.go                        |  2 +-
 pkg/api/staking_test.go                   |  3 +-
 pkg/storageincentives/staking/contract.go | 48 +++++++++++++++--------
 4 files changed, 35 insertions(+), 19 deletions(-)

diff --git a/pkg/api/export_test.go b/pkg/api/export_test.go
index f5d6fe16ec0..75fba7c4e24 100644
--- a/pkg/api/export_test.go
+++ b/pkg/api/export_test.go
@@ -100,6 +100,7 @@ type (
 	BucketData                        = bucketData
 	WalletResponse                    = walletResponse
 	GetStakeResponse                  = getStakeResponse
+	WithdrawAllStakeResponse          = withdrawAllStakeResponse
 )
 
 var (
diff --git a/pkg/api/staking.go b/pkg/api/staking.go
index 8252416f2a8..77e8a7e836c 100644
--- a/pkg/api/staking.go
+++ b/pkg/api/staking.go
@@ -97,7 +97,7 @@ func (s *Service) getStakedAmountHandler(w http.ResponseWriter, r *http.Request)
 }
 
 func (s *Service) withdrawAllStakeHandler(w http.ResponseWriter, r *http.Request) {
-	logger := s.logger.WithName("withdraw_all_stake").Build()
+	logger := s.logger.WithName("delete_withdraw_all_stake").Build()
 
 	txHash, err := s.stakingContract.WithdrawAllStake(r.Context())
 	if err != nil {
diff --git a/pkg/api/staking_test.go b/pkg/api/staking_test.go
index eb4a39dfead..2ca045cbcaa 100644
--- a/pkg/api/staking_test.go
+++ b/pkg/api/staking_test.go
@@ -186,7 +186,8 @@ func TestWithdrawAllStake(t *testing.T) {
 			}),
 		)
 		ts, _, _, _ := newTestServer(t, testServerOptions{DebugAPI: true, StakingContract: contract})
-		jsonhttptest.Request(t, ts, http.MethodDelete, "/stake", http.StatusOK)
+		jsonhttptest.Request(t, ts, http.MethodDelete, "/stake", http.StatusOK, jsonhttptest.WithExpectedJSONResponse(
+			&api.WithdrawAllStakeResponse{TxHash: txHash.String()}))
 	})
 
 	t.Run("with invalid stake amount", func(t *testing.T) {
diff --git a/pkg/storageincentives/staking/contract.go b/pkg/storageincentives/staking/contract.go
index cce2838e868..fcf99935277 100644
--- a/pkg/storageincentives/staking/contract.go
+++ b/pkg/storageincentives/staking/contract.go
@@ -162,42 +162,46 @@ func (c *contract) getStake(ctx context.Context, overlay swarm.Address) (*big.In
 	if err != nil {
 		return nil, err
 	}
+
+	if len(results) == 0 {
+		return nil, errors.New("unexpected empty results")
+	}
+
 	return abi.ConvertType(results[0], new(big.Int)).(*big.Int), nil
 }
 
-func (c *contract) DepositStake(ctx context.Context, stakedAmount *big.Int) (txHash common.Hash, err error) {
+func (c *contract) DepositStake(ctx context.Context, stakedAmount *big.Int) (common.Hash, error) {
 	prevStakedAmount, err := c.GetStake(ctx)
 	if err != nil {
-		return
+		return common.Hash{}, err
 	}
 
 	if len(prevStakedAmount.Bits()) == 0 {
 		if stakedAmount.Cmp(MinimumStakeAmount) == -1 {
-			err = ErrInsufficientStakeAmount
-			return
+			return common.Hash{}, ErrInsufficientStakeAmount
 		}
 	}
 
 	balance, err := c.getBalance(ctx)
 	if err != nil {
-		return
+		return common.Hash{}, err
 	}
 
 	if balance.Cmp(stakedAmount) < 0 {
-		err = ErrInsufficientFunds
-		return
+		return common.Hash{}, ErrInsufficientFunds
 	}
 
 	_, err = c.sendApproveTransaction(ctx, stakedAmount)
 	if err != nil {
-		return
+		return common.Hash{}, err
 	}
 
 	receipt, err := c.sendDepositStakeTransaction(ctx, c.owner, stakedAmount, c.overlayNonce)
-	if receipt != nil {
-		txHash = receipt.TxHash
+	if err != nil {
+		return common.Hash{}, err
 	}
-	return
+
+	return receipt.TxHash, nil
 }
 
 func (c *contract) GetStake(ctx context.Context) (*big.Int, error) {
@@ -226,6 +230,11 @@ func (c *contract) getBalance(ctx context.Context) (*big.Int, error) {
 	if err != nil {
 		return nil, err
 	}
+
+	if len(results) == 0 {
+		return nil, errors.New("unexpected empty results")
+	}
+
 	return abi.ConvertType(results[0], new(big.Int)).(*big.Int), nil
 }
 
@@ -235,8 +244,7 @@ func (c *contract) WithdrawAllStake(ctx context.Context) (txHash common.Hash, er
 		return
 	}
 	if !isPaused {
-		err = ErrNotPaused
-		return
+		return common.Hash{}, ErrNotPaused
 	}
 
 	stakedAmount, err := c.getStake(ctx, c.overlay)
@@ -245,20 +253,22 @@ func (c *contract) WithdrawAllStake(ctx context.Context) (txHash common.Hash, er
 	}
 
 	if stakedAmount.Cmp(big.NewInt(0)) <= 0 {
-		err = ErrInsufficientStake
-		return
+		return common.Hash{}, ErrInsufficientStake
 	}
 
 	_, err = c.sendApproveTransaction(ctx, stakedAmount)
 	if err != nil {
-		return
+		return common.Hash{}, err
 	}
 
 	receipt, err := c.withdrawFromStake(ctx, stakedAmount)
+	if err != nil {
+		return common.Hash{}, err
+	}
 	if receipt != nil {
 		txHash = receipt.TxHash
 	}
-	return
+	return txHash, nil
 }
 
 func (c *contract) withdrawFromStake(ctx context.Context, stakedAmount *big.Int) (*types.Receipt, error) {
@@ -297,5 +307,9 @@ func (c *contract) paused(ctx context.Context) (bool, error) {
 		return false, err
 	}
 
+	if len(results) == 0 {
+		return false, errors.New("unexpected empty results")
+	}
+
 	return results[0].(bool), nil
 }