Merge pull request #590 from c9s/fix-net-asset

fix: fix net asset field

pkg/exchange/max/exchange.go

@@ -663,10 +663,12 @@ func (e *Exchange) QueryAccount(ctx context.Context) (*types.Account, error) {
 
 	var balances = make(types.BalanceMap)
 	for _, a := range userInfo.Accounts {
+		cur := toGlobalCurrency(a.Currency)
 		balances[toGlobalCurrency(a.Currency)] = types.Balance{
-			Currency:  toGlobalCurrency(a.Currency),
+			Currency:  cur,
 			Available: a.Balance,
 			Locked:    a.Locked,
+			NetAsset:  a.Balance.Add(a.Locked),
 		}
 	}
 
@@ -861,10 +863,12 @@ func (e *Exchange) QueryAccountBalances(ctx context.Context) (types.BalanceMap,
 	var balances = make(types.BalanceMap)
 
 	for _, a := range accounts {
-		balances[toGlobalCurrency(a.Currency)] = types.Balance{
+		cur := toGlobalCurrency(a.Currency)
-			Currency:  toGlobalCurrency(a.Currency),
+		balances[cur] = types.Balance{
+			Currency:  cur,
 			Available: a.Balance,
 			Locked:    a.Locked,
+			NetAsset:  a.Balance.Add(a.Locked),
 		}
 	}
 

pkg/types/account_test.go

@@ -8,36 +8,6 @@ import (
 	"github.com/c9s/bbgo/pkg/fixedpoint"
 )
 
-func TestBalanceMap_Add(t *testing.T) {
-	var bm = BalanceMap{}
-	var bm2 = bm.Add(BalanceMap{
-		"BTC": Balance{
-			Currency:  "BTC",
-			Available: fixedpoint.MustNewFromString("10.0"),
-			Locked:    fixedpoint.MustNewFromString("0"),
-			NetAsset:  fixedpoint.MustNewFromString("10.0"),
-		},
-	})
-	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) {
 	a := NewAccount()
 	a.AddBalance("USDT", fixedpoint.NewFromInt(1000))

pkg/types/asset.go

@@ -12,10 +12,13 @@ import (
 
 type Asset struct {
 	Currency string           `json:"currency" db:"currency"`
+
 	Total    fixedpoint.Value `json:"total" db:"total"`
 
 	NetAsset fixedpoint.Value `json:"netAsset" db:"net_asset"`
 
+	Interest fixedpoint.Value `json:"interest" db:"interest"`
+
 	// InUSD is net asset in USD
 	InUSD fixedpoint.Value `json:"inUSD" db:"net_asset_in_usd"`
 

pkg/types/balance.go

@@ -23,6 +23,16 @@ type Balance struct {
 	NetAsset fixedpoint.Value `json:"net,omitempty"`
 }
 
+func (b Balance) Add(b2 Balance) Balance {
+	var newB = b
+	newB.Available = b.Available.Add(b2.Available)
+	newB.Locked = b.Locked.Add(b2.Locked)
+	newB.Borrowed = b.Borrowed.Add(b2.Borrowed)
+	newB.NetAsset = b.NetAsset.Add(b2.NetAsset)
+	newB.Interest = b.Interest.Add(b2.Interest)
+	return newB
+}
+
 func (b Balance) Total() fixedpoint.Value {
 	return b.Available.Add(b.Locked)
 }
@@ -65,11 +75,7 @@ func (m BalanceMap) Add(bm BalanceMap) BalanceMap {
 	for _, b := range bm {
 		tb, ok := total[b.Currency]
 		if ok {
-			tb.Available = tb.Available.Add(b.Available)
+			tb = tb.Add(b)
-			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
 		}
@@ -102,12 +108,11 @@ func (m BalanceMap) Assets(prices map[string]fixedpoint.Value, priceTime time.Ti
 	_, btcInUSD, hasBtcPrice := findUSDMarketPrice("BTC", prices)
 
 	for currency, b := range m {
-		if b.Locked.IsZero() && b.Available.IsZero() && b.Borrowed.IsZero() {
-			continue
-		}
-
 		total := b.Total()
 		netAsset := b.Net()
+		if total.IsZero() && netAsset.IsZero() {
+			continue
+		}
 
 		asset := Asset{
 			Currency:  currency,
@@ -116,6 +121,7 @@ func (m BalanceMap) Assets(prices map[string]fixedpoint.Value, priceTime time.Ti
 			Locked:    b.Locked,
 			Available: b.Available,
 			Borrowed:  b.Borrowed,
+			Interest:  b.Interest,
 			NetAsset:  netAsset,
 		}
 

pkg/types/balance_test.go

@@ -0,0 +1,39 @@
+package types
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+
+	"github.com/c9s/bbgo/pkg/fixedpoint"
+)
+
+func TestBalanceMap_Add(t *testing.T) {
+	var bm = BalanceMap{}
+	var bm2 = bm.Add(BalanceMap{
+		"BTC": Balance{
+			Currency:  "BTC",
+			Available: fixedpoint.MustNewFromString("10.0"),
+			Locked:    fixedpoint.MustNewFromString("0"),
+			NetAsset:  fixedpoint.MustNewFromString("10.0"),
+		},
+	})
+	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)
+}