Merge pull request #71 from Zondax/balance_fix

Fix /account to match updated requirements
Zondax · Oct 2, 2020 · b809c96 · b809c96
2 parents 00a958b + ed1d225
commit b809c96
Show file tree

Hide file tree

Showing 3 changed files with 105 additions and 10 deletions.
diff --git 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
@@ -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
@@ -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")
+	}
+}