From 4404098bf92c403bcafc6920e458c63cec3cdbc9 Mon Sep 17 00:00:00 2001
From: c9s <yoanlin93@gmail.com>
Date: Wed, 4 May 2022 17:39:35 +0800
Subject: [PATCH] fix balance map add

---
 pkg/types/account_test.go | 17 +++++++++++++++++
 pkg/types/balance.go      | 24 +++++++++++++-----------
 2 files changed, 30 insertions(+), 11 deletions(-)

diff --git a/pkg/types/account_test.go b/pkg/types/account_test.go
index b5c1e0b3b..aa16f73c7 100644
--- a/pkg/types/account_test.go
+++ b/pkg/types/account_test.go
@@ -19,6 +19,23 @@ func TestBalanceMap_Add(t *testing.T) {
 		},
 	})
 	assert.Len(t, bm2, 1)
+
+	var bm3 = bm2.Add(BalanceMap{
+		"BTC": Balance{
+			Currency:  "BTC",
+			Available: fixedpoint.MustNewFromString("1.0"),
+			Locked:    fixedpoint.MustNewFromString("0"),
+			NetAsset:  fixedpoint.MustNewFromString("1.0"),
+		},
+		"LTC": Balance{
+			Currency:  "LTC",
+			Available: fixedpoint.MustNewFromString("20.0"),
+			Locked:    fixedpoint.MustNewFromString("0"),
+			NetAsset:  fixedpoint.MustNewFromString("20.0"),
+		},
+	})
+	assert.Len(t, bm3, 2)
+	assert.Equal(t, fixedpoint.MustNewFromString("11.0"), bm3["BTC"].Available)
 }
 
 func TestAccountLockAndUnlock(t *testing.T) {
diff --git a/pkg/types/balance.go b/pkg/types/balance.go
index da2a94923..fa17ae892 100644
--- a/pkg/types/balance.go
+++ b/pkg/types/balance.go
@@ -41,7 +41,6 @@ func (b Balance) String() (o string) {
 	return o
 }
 
-
 type BalanceMap map[string]Balance
 
 func (m BalanceMap) Currencies() (currencies []string) {
@@ -52,14 +51,18 @@ func (m BalanceMap) Currencies() (currencies []string) {
 }
 
 func (m BalanceMap) Add(bm BalanceMap) BalanceMap {
-	var total = BalanceMap{}
+	var total = m.Copy()
 	for _, b := range bm {
-		tb := total[b.Currency]
-		tb.Available = tb.Available.Add(b.Available)
-		tb.Locked = tb.Locked.Add(b.Locked)
-		tb.Borrowed = tb.Borrowed.Add(b.Borrowed)
-		tb.NetAsset = tb.NetAsset.Add(b.NetAsset)
-		tb.Interest = tb.Interest.Add(b.Interest)
+		tb, ok := total[b.Currency]
+		if ok {
+			tb.Available = tb.Available.Add(b.Available)
+			tb.Locked = tb.Locked.Add(b.Locked)
+			tb.Borrowed = tb.Borrowed.Add(b.Borrowed)
+			tb.NetAsset = tb.NetAsset.Add(b.NetAsset)
+			tb.Interest = tb.Interest.Add(b.Interest)
+		} else {
+			tb = b
+		}
 		total[b.Currency] = tb
 	}
 	return total
@@ -107,14 +110,14 @@ func (m BalanceMap) Assets(prices map[string]fixedpoint.Value, priceTime time.Ti
 			NetAsset:  netAsset,
 		}
 
-		if strings.HasPrefix(currency,"USD") { // for usd
+		if strings.HasPrefix(currency, "USD") { // for usd
 			asset.InUSD = netAsset
 			asset.PriceInUSD = fixedpoint.One
 			if hasBtcPrice && !asset.InUSD.IsZero() {
 				asset.InBTC = asset.InUSD.Div(btcusdt)
 			}
 		} else { // for crypto
-			if market, usdPrice, ok := findUSDMarketPrice(currency, prices) ; ok {
+			if market, usdPrice, ok := findUSDMarketPrice(currency, prices); ok {
 				// this includes USDT, USD, USDC and so on
 				if strings.HasPrefix(market, "USD") {
 					if !asset.Total.IsZero() {
@@ -154,7 +157,6 @@ func (m BalanceMap) Print() {
 	}
 }
 
-
 func findUSDMarketPrice(currency string, prices map[string]fixedpoint.Value) (string, fixedpoint.Value, bool) {
 	usdMarkets := []string{currency + "USDT", currency + "USDC", currency + "USD", "USDT" + currency}
 	for _, market := range usdMarkets {