From 271e15aae8e251acc07fc1199a592456cbc540ce Mon Sep 17 00:00:00 2001
From: Calin Martinconi <martinconic@gmail.com>
Date: Wed, 4 Sep 2024 18:29:28 +0300
Subject: [PATCH] fix: rename only withdrawable stake response (#414)

* fix: rename only withdrawable stake response

* fix: use and add the right path to withdrawable stake

* fix: use of  withdrawable stake

* fix: expected withdrawable stake error
---
 pkg/bee/api/stake.go     | 22 ++++++++++++++++++----
 pkg/bee/client.go        |  5 +++++
 pkg/check/stake/stake.go | 36 +++++++++++++++++++++++++++++++++---
 3 files changed, 56 insertions(+), 7 deletions(-)

diff --git a/pkg/bee/api/stake.go b/pkg/bee/api/stake.go
index e5124c14..3d548e8f 100644
--- a/pkg/bee/api/stake.go
+++ b/pkg/bee/api/stake.go
@@ -13,7 +13,11 @@ import (
 type StakingService service
 
 type getStakeResponse struct {
-	WithdrawableStake *bigint.BigInt `json:"withdrawableStake"`
+	StakedAmount *bigint.BigInt `json:"stakedAmount"`
+}
+
+type getWithdrawableResponse struct {
+	WithdrawableAmount *bigint.BigInt `json:"withdrawableAmount"`
 }
 type stakeDepositResponse struct {
 	TxHash string `json:"txhash"`
@@ -32,14 +36,24 @@ func (s *StakingService) DepositStake(ctx context.Context, amount *big.Int) (txH
 	return r.TxHash, nil
 }
 
-// GetWithdrawableStake gets stake
-func (s *StakingService) GetWithdrawableStake(ctx context.Context) (withdrawableStake *big.Int, err error) {
+// GetStakedAmount gets stake
+func (s *StakingService) GetStakedAmount(ctx context.Context) (stakedAmount *big.Int, err error) {
 	r := new(getStakeResponse)
 	err = s.client.requestJSON(ctx, http.MethodGet, "/stake", nil, r)
 	if err != nil {
 		return nil, err
 	}
-	return r.WithdrawableStake.Int, nil
+	return r.StakedAmount.Int, nil
+}
+
+// GetWithdrawableStake gets stake
+func (s *StakingService) GetWithdrawableStake(ctx context.Context) (withdrawableStake *big.Int, err error) {
+	r := new(getWithdrawableResponse)
+	err = s.client.requestJSON(ctx, http.MethodGet, "/stake/withdrawable", nil, r)
+	if err != nil {
+		return nil, err
+	}
+	return r.WithdrawableAmount.Int, nil
 }
 
 // MigrateStake withdraws stake
diff --git a/pkg/bee/client.go b/pkg/bee/client.go
index 1468d408..e362c11e 100644
--- a/pkg/bee/client.go
+++ b/pkg/bee/client.go
@@ -884,6 +884,11 @@ func (c *Client) DepositStake(ctx context.Context, amount *big.Int) (string, err
 
 // GetStake returns stake amount
 func (c *Client) GetStake(ctx context.Context) (*big.Int, error) {
+	return c.api.Stake.GetStakedAmount(ctx)
+}
+
+// GetWithdrawableStake returns withdrawable amount
+func (c *Client) GetWithdrawableStake(ctx context.Context) (*big.Int, error) {
 	return c.api.Stake.GetWithdrawableStake(ctx)
 }
 
diff --git a/pkg/check/stake/stake.go b/pkg/check/stake/stake.go
index 69af4a8c..df83a0d5 100644
--- a/pkg/check/stake/stake.go
+++ b/pkg/check/stake/stake.go
@@ -109,15 +109,23 @@ func (c *Check) Run(ctx context.Context, cluster orchestration.Cluster, opts int
 		return err
 	}
 
+	if err := expectWithdrawableStake(ctx, client, o.Amount); err != nil {
+		return err
+	}
+
 	// should allow increasing the stake amount
-	withdrawableStake := new(big.Int).Add(o.Amount, big.NewInt(1))
+	stakedAmount := new(big.Int).Add(o.Amount, big.NewInt(1))
 
 	_, err = client.DepositStake(ctx, big.NewInt(1))
 	if err != nil {
 		return fmt.Errorf("increase stake amount: %w", err)
 	}
 
-	if err := expectStakeAmountIs(ctx, client, withdrawableStake); err != nil {
+	if err := expectStakeAmountIs(ctx, client, stakedAmount); err != nil {
+		return err
+	}
+
+	if err := expectWithdrawableStake(ctx, client, stakedAmount); err != nil {
 		return err
 	}
 
@@ -127,7 +135,7 @@ func (c *Check) Run(ctx context.Context, cluster orchestration.Cluster, opts int
 		return errors.New("withdraw from running contract should fail")
 	}
 
-	if err := expectStakeAmountIs(ctx, client, withdrawableStake); err != nil {
+	if err := expectStakeAmountIs(ctx, client, stakedAmount); err != nil {
 		return err
 	}
 
@@ -157,6 +165,10 @@ func (c *Check) Run(ctx context.Context, cluster orchestration.Cluster, opts int
 		return err
 	}
 
+	if err := expectWithdrawableStake(ctx, client, zero); err != nil {
+		return err
+	}
+
 	return nil
 }
 
@@ -172,3 +184,21 @@ func expectStakeAmountIs(ctx context.Context, client *bee.Client, expected *big.
 
 	return nil
 }
+
+func expectWithdrawableStake(ctx context.Context, client *bee.Client, expected *big.Int) error {
+	withdrawable, err := client.GetWithdrawableStake(ctx)
+	if err != nil {
+		return fmt.Errorf("get stake amount: %w", err)
+	}
+
+	if (expected.Cmp(zero) == 0 && withdrawable.Cmp(expected) != 0) ||
+		(expected.Cmp(zero) != 0 && withdrawable.Cmp(zero) == 0) {
+		if expected.Cmp(zero) == 0 {
+			return fmt.Errorf("expected withdrawable stake to be %d, got: %d", expected, withdrawable)
+		} else {
+			return fmt.Errorf("expected withdrawable stake should not be equal to 0")
+		}
+	}
+
+	return nil
+}