From ed1d225c9983150fc109ae801fceb5c1810d6ad2 Mon Sep 17 00:00:00 2001
From: ER <raynaudo.ee@gmail.com>
Date: Fri, 2 Oct 2020 12:33:58 -0300
Subject: [PATCH] Fix balance to match updated requirements Add GetBalance test

---
 rosetta/services/account.go   | 18 +++----
 rosetta/services/constants.go |  1 +
 rosetta/tests/rosetta_test.go | 96 +++++++++++++++++++++++++++++++++++
 3 files changed, 105 insertions(+), 10 deletions(-)

diff --git a/rosetta/services/account.go b/rosetta/services/account.go
index a67ae5d..711d484 100644
--- a/rosetta/services/account.go
+++ b/rosetta/services/account.go
@@ -110,6 +110,12 @@ func (a AccountAPIService) AccountBalance(ctx context.Context,
 		case LockedBalanceStr:
 			lockedFunds := stateMultisig.AmountLocked(queryTipSet.Height())
 			balanceStr = lockedFunds.String()
+		case SpendableBalanceStr:
+			available, err := a.node.MsigGetAvailableBalance(ctx, addr, queryTipSet.Key())
+			if err != nil {
+				return nil, BuildError(ErrUnableToGetBalance, err)
+			}
+			balanceStr = available.String()
 		case VestingScheduleStr:
 			stEpoch := stateMultisig.StartEpoch.String()
 			unlockDuration := stateMultisig.UnlockDuration.String()
@@ -121,16 +127,8 @@ func (a AccountAPIService) AccountBalance(ctx context.Context,
 			return nil, BuildError(ErrMustSpecifySubAccount, nil)
 		}
 	} else {
-		//Get available balance
-		if isMultiSig {
-			balance, err := a.node.MsigGetAvailableBalance(ctx, addr, queryTipSet.Key())
-			if err != nil {
-				return nil, BuildError(ErrUnableToGetBalance, err)
-			}
-			balanceStr = balance.String()
-		} else {
-			balanceStr = actor.Balance.String()
-		}
+		//Get available balance (spendable + locked)
+		balanceStr = actor.Balance.String()
 	}
 
 	queryTipSetHeight := int64(queryTipSet.Height())
diff --git a/rosetta/services/constants.go b/rosetta/services/constants.go
index 5c939d0..2218192 100644
--- a/rosetta/services/constants.go
+++ b/rosetta/services/constants.go
@@ -31,6 +31,7 @@ const (
 
 	///Account
 	LockedBalanceStr         = "LockedBalance"
+	SpendableBalanceStr      = "SpendableBalance"
 	VestingScheduleStr       = "VestingSchedule"
 	LockedFundsKey           = "LockedFunds"
 	VestingStartEpochKey     = "StartEpoch"
diff --git a/rosetta/tests/rosetta_test.go b/rosetta/tests/rosetta_test.go
index 7009526..c67ab82 100644
--- a/rosetta/tests/rosetta_test.go
+++ b/rosetta/tests/rosetta_test.go
@@ -391,3 +391,99 @@ func TestSendTransaction(t *testing.T) {
 		t.Fatal("NOT MATCHING")
 	}
 }
+
+func TestGetBalance(t *testing.T) {
+
+	rosettaClient := setupRosettaClient()
+	testAddress := "t023232"
+
+	fmt.Println("Testing on address:", testAddress)
+
+	// Get full balance (locked + spendable)
+	req := &types.AccountBalanceRequest{
+		NetworkIdentifier: NetworkID,
+		AccountIdentifier: &types.AccountIdentifier{
+			Address:    testAddress,
+			SubAccount: nil,
+		},
+	}
+
+	resp, err1, err2 := rosettaClient.AccountAPI.AccountBalance(ctx, req)
+	if err1 != nil {
+		t.Fatal(err1.Message)
+	}
+
+	if err2 != nil {
+		t.Fatal(err2.Error())
+	}
+
+	if resp == nil {
+		t.Fatal()
+	}
+
+	full := resp.Balances[0].Value
+	fmt.Println("Full balance is:", full)
+
+	// Get locked balance
+	req = &types.AccountBalanceRequest{
+		NetworkIdentifier: NetworkID,
+		AccountIdentifier: &types.AccountIdentifier{
+			Address: testAddress,
+			SubAccount: &types.SubAccountIdentifier{
+				Address: services.LockedBalanceStr,
+			},
+		},
+	}
+
+	resp, err1, err2 = rosettaClient.AccountAPI.AccountBalance(ctx, req)
+	if err1 != nil {
+		t.Fatal(err1.Message)
+	}
+
+	if err2 != nil {
+		t.Fatal(err2.Error())
+	}
+
+	if resp == nil {
+		t.Fatal()
+	}
+
+	locked := resp.Balances[0].Value
+	fmt.Println("Locked balance is:", locked)
+
+	// Get spendable balance
+	req = &types.AccountBalanceRequest{
+		NetworkIdentifier: NetworkID,
+		AccountIdentifier: &types.AccountIdentifier{
+			Address: testAddress,
+			SubAccount: &types.SubAccountIdentifier{
+				Address: services.SpendableBalanceStr,
+			},
+		},
+	}
+
+	resp, err1, err2 = rosettaClient.AccountAPI.AccountBalance(ctx, req)
+	if err1 != nil {
+		t.Fatal(err1.Message)
+	}
+
+	if err2 != nil {
+		t.Fatal(err2.Error())
+	}
+
+	if resp == nil {
+		t.Fatal()
+	}
+
+	spendable := resp.Balances[0].Value
+	fmt.Println("Spendable balance is:", spendable)
+
+	//Check that values complies
+	nfull, _ := strconv.Atoi(full)
+	nlocked, _ := strconv.Atoi(locked)
+	nspendable, _ := strconv.Atoi(spendable)
+
+	if (nlocked + nspendable) != nfull {
+		t.Fatal("Balances mismatch")
+	}
+}