risk: add test case for account calculator

pkg/risk/account_value.go

@@ -38,6 +38,10 @@ func (c *AccountValueCalculator) UpdatePrices(ctx context.Context) error {
 	currencies := balances.Currencies()
 	var symbols []string
 	for _, currency := range currencies {
+		if currency == c.quoteCurrency {
+			continue
+		}
+
 		symbol := currency + c.quoteCurrency
 		symbols = append(symbols, symbol)
 	}
@@ -80,6 +84,34 @@ func (c *AccountValueCalculator) DebtValue(ctx context.Context) (fixedpoint.Valu
 	return debtValue, nil
 }
 
+func (c *AccountValueCalculator) MarketValue(ctx context.Context) (fixedpoint.Value, error) {
+	marketValue := fixedpoint.Zero
+
+	if len(c.prices) == 0 {
+		if err := c.UpdatePrices(ctx); err != nil {
+			return marketValue, err
+		}
+	}
+
+	balances := c.session.Account.Balances()
+	for _, b := range balances {
+		if b.Currency == c.quoteCurrency {
+			marketValue = marketValue.Add(b.Total())
+			continue
+		}
+
+		symbol := b.Currency + c.quoteCurrency
+		price, ok := c.prices[symbol]
+		if !ok {
+			continue
+		}
+
+		marketValue = marketValue.Add(b.Total().Mul(price))
+	}
+
+	return marketValue, nil
+}
+
 func (c *AccountValueCalculator) NetValue(ctx context.Context) (fixedpoint.Value, error) {
 	accountValue := fixedpoint.Zero
 
@@ -91,6 +123,11 @@ func (c *AccountValueCalculator) NetValue(ctx context.Context) (fixedpoint.Value
 
 	balances := c.session.Account.Balances()
 	for _, b := range balances {
+		if b.Currency == c.quoteCurrency {
+			accountValue = accountValue.Add(b.Net())
+			continue
+		}
+
 		symbol := b.Currency + c.quoteCurrency
 		price, ok := c.prices[symbol]
 		if !ok {
@@ -103,17 +140,6 @@ func (c *AccountValueCalculator) NetValue(ctx context.Context) (fixedpoint.Value
 	return accountValue, nil
 }
 
-func CalculateAccountNetValue(session *bbgo.ExchangeSession) (fixedpoint.Value, error) {
-	accountValue := fixedpoint.Zero
-	ctx := context.Background()
-	c := NewAccountValueCalculator(session, "USDT")
-	if err := c.UpdatePrices(ctx); err != nil {
-		return accountValue, err
-	}
-
-	return c.NetValue(ctx)
-}
-
 func CalculateBaseQuantity(session *bbgo.ExchangeSession, market types.Market, price, quantity, leverage fixedpoint.Value) (fixedpoint.Value, error) {
 	// default leverage guard
 	if leverage.IsZero() {

pkg/risk/account_value_test.go

@@ -0,0 +1,113 @@
+package risk
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/golang/mock/gomock"
+	"github.com/stretchr/testify/assert"
+
+	"github.com/c9s/bbgo/pkg/bbgo"
+	"github.com/c9s/bbgo/pkg/fixedpoint"
+	"github.com/c9s/bbgo/pkg/types"
+	"github.com/c9s/bbgo/pkg/types/mocks"
+)
+
+func newTestTicker() types.Ticker {
+	return types.Ticker{
+		Time:   time.Now(),
+		Volume: fixedpoint.Zero,
+		Last:   fixedpoint.NewFromFloat(19000.0),
+		Open:   fixedpoint.NewFromFloat(19500.0),
+		High:   fixedpoint.NewFromFloat(19900.0),
+		Low:    fixedpoint.NewFromFloat(18800.0),
+		Buy:    fixedpoint.NewFromFloat(19500.0),
+		Sell:   fixedpoint.NewFromFloat(18900.0),
+	}
+}
+
+func TestAccountValueCalculator_NetValue(t *testing.T) {
+
+	t.Run("borrow and available", func(t *testing.T) {
+		mockCtrl := gomock.NewController(t)
+		defer mockCtrl.Finish()
+
+		mockEx := mocks.NewMockExchange(mockCtrl)
+		// for market data stream and user data stream
+		mockEx.EXPECT().NewStream().Return(&types.StandardStream{}).Times(2)
+		mockEx.EXPECT().QueryTickers(gomock.Any(), []string{"BTCUSDT"}).Return(map[string]types.Ticker{
+			"BTCUSDT": newTestTicker(),
+		}, nil)
+
+		session := bbgo.NewExchangeSession("test", mockEx)
+		session.Account.UpdateBalances(types.BalanceMap{
+			"BTC": {
+				Currency:  "BTC",
+				Available: fixedpoint.NewFromFloat(2.0),
+				Locked:    fixedpoint.Zero,
+				Borrowed:  fixedpoint.NewFromFloat(1.0),
+				Interest:  fixedpoint.Zero,
+				NetAsset:  fixedpoint.Zero,
+			},
+			"USDT": {
+				Currency:  "USDT",
+				Available: fixedpoint.NewFromFloat(1000.0),
+				Locked:    fixedpoint.Zero,
+				Borrowed:  fixedpoint.Zero,
+				Interest:  fixedpoint.Zero,
+				NetAsset:  fixedpoint.Zero,
+			},
+		})
+		assert.NotNil(t, session)
+
+		cal := NewAccountValueCalculator(session, "USDT")
+		assert.NotNil(t, cal)
+
+		ctx := context.Background()
+		netValue, err := cal.NetValue(ctx)
+		assert.NoError(t, err)
+		assert.Equal(t, "20000", netValue.String())
+	})
+
+	t.Run("borrowed and sold", func(t *testing.T) {
+		mockCtrl := gomock.NewController(t)
+		defer mockCtrl.Finish()
+
+		mockEx := mocks.NewMockExchange(mockCtrl)
+		// for market data stream and user data stream
+		mockEx.EXPECT().NewStream().Return(&types.StandardStream{}).Times(2)
+		mockEx.EXPECT().QueryTickers(gomock.Any(), []string{"BTCUSDT"}).Return(map[string]types.Ticker{
+			"BTCUSDT": newTestTicker(),
+		}, nil)
+
+		session := bbgo.NewExchangeSession("test", mockEx)
+		session.Account.UpdateBalances(types.BalanceMap{
+			"BTC": {
+				Currency:  "BTC",
+				Available: fixedpoint.Zero,
+				Locked:    fixedpoint.Zero,
+				Borrowed:  fixedpoint.NewFromFloat(1.0),
+				Interest:  fixedpoint.Zero,
+				NetAsset:  fixedpoint.Zero,
+			},
+			"USDT": {
+				Currency:  "USDT",
+				Available: fixedpoint.NewFromFloat(21000.0),
+				Locked:    fixedpoint.Zero,
+				Borrowed:  fixedpoint.Zero,
+				Interest:  fixedpoint.Zero,
+				NetAsset:  fixedpoint.Zero,
+			},
+		})
+		assert.NotNil(t, session)
+
+		cal := NewAccountValueCalculator(session, "USDT")
+		assert.NotNil(t, cal)
+
+		ctx := context.Background()
+		netValue, err := cal.NetValue(ctx)
+		assert.NoError(t, err)
+		assert.Equal(t, "2000", netValue.String()) // 21000-19000
+	})
+}