fix stock calculation

bbgo/stock.go

@@ -8,6 +8,14 @@ import (
 	"strings"
 )
 
+func zero(a float64) bool {
+	return int(math.Round(a*1e8)) == 0
+}
+
+func round(a float64) float64 {
+	return math.Round(a*1e8) / 1e8
+}
+
 type Stock types.Trade
 
 func (stock *Stock) String() string {
@@ -15,9 +23,9 @@ func (stock *Stock) String() string {
 }
 
 func (stock *Stock) Consume(quantity float64) float64 {
-	delta := math.Min(stock.Quantity, quantity)
-	stock.Quantity -= delta
-	return delta
+	q := math.Min(stock.Quantity, quantity)
+	stock.Quantity = round(stock.Quantity - q)
+	return q
 }
 
 type StockSlice []Stock
@@ -27,7 +35,7 @@ func (slice StockSlice) Quantity() (total float64) {
 		total += stock.Quantity
 	}
 
-	return total
+	return round(total)
 }
 
 type StockManager struct {
@@ -39,10 +47,15 @@ type StockManager struct {
 
 func (m *StockManager) Stock(buy Stock) error {
 	m.Stocks = append(m.Stocks, buy)
+	return m.flushPendingSells()
+}
+
+func (m *StockManager) flushPendingSells() error {
+	if len(m.Stocks) > 0 && len(m.PendingSells) > 0 {
 
-	if len(m.PendingSells) > 0 {
 		pendingSells := m.PendingSells
 		m.PendingSells = nil
+
 		for _, sell := range pendingSells {
 			if err := m.Consume(sell); err != nil {
 				return err
@@ -68,14 +81,15 @@ func (m *StockManager) Consume(sell Stock) error {
 			continue
 		}
 
-		sell.Quantity -= stock.Consume(sell.Quantity)
-		m.Stocks[idx] = stock
-
-		if math.Round(stock.Quantity*1e8) == 0.0 {
-			m.Stocks = m.Stocks[:idx]
+		if zero(stock.Quantity) {
+			continue
 		}
 
-		if math.Round(sell.Quantity*1e8) == 0.0 {
+		delta := stock.Consume(sell.Quantity)
+		sell.Consume(delta)
+		m.Stocks[idx] = stock
+
+		if zero(sell.Quantity) {
 			return nil
 		}
 	}
@@ -83,16 +97,18 @@ func (m *StockManager) Consume(sell Stock) error {
 	idx = len(m.Stocks) - 1
 	for ; idx >= 0; idx-- {
 		stock := m.Stocks[idx]
-		sell.Quantity -= stock.Consume(sell.Quantity)
-		m.Stocks[idx] = stock
 
-		if math.Round(stock.Quantity*1e8) == 0.0 {
-			// remove the latest stock
-			m.Stocks = m.Stocks[:idx]
+		if zero(stock.Quantity) {
+			continue
 		}
+
+		delta := stock.Consume(sell.Quantity)
+		sell.Consume(delta)
+
+		m.Stocks[idx] = stock
 	}
 
-	if math.Round(sell.Quantity*1e8) > 0.0 {
+	if !zero(sell.Quantity) {
 		m.PendingSells = append(m.PendingSells, sell)
 	}
 
@@ -114,26 +130,29 @@ func (m *StockManager) LoadTrades(trades []types.Trade) (checkpoints []int, err
 			}
 		}
 
-		if trade.Symbol == m.Symbol {
-			if trade.IsBuyer {
-				if idx > 0 && len(m.Stocks) == 0 {
-					checkpoints = append(checkpoints, idx)
-				}
+		if trade.Symbol != m.Symbol {
+			continue
+		}
 
-				stock := toStock(trade)
-				if err := m.Stock(stock); err != nil {
-					return checkpoints, err
-				}
-			} else {
-				stock := toStock(trade)
-				if err := m.Consume(stock); err != nil {
-					return checkpoints, err
-				}
+		if trade.IsBuyer {
+			if idx > 0 && len(m.Stocks) == 0 {
+				checkpoints = append(checkpoints, idx)
+			}
+
+			stock := toStock(trade)
+			if err := m.Stock(stock); err != nil {
+				return checkpoints, err
+			}
+		} else {
+			stock := toStock(trade)
+			if err := m.Consume(stock); err != nil {
+				return checkpoints, err
 			}
 		}
 	}
 
-	return checkpoints, nil
+	err = m.flushPendingSells()
+	return checkpoints, err
 }
 
 func toStock(trade types.Trade) Stock {
@@ -143,7 +162,7 @@ func toStock(trade types.Trade) Stock {
 		} else {
 			trade.Quantity += trade.Fee
 		}
-		trade.Fee = 0
+		trade.Fee = 0.0
 	}
 	return Stock(trade)
 }

bbgo/stock_test.go

@@ -1,13 +1,37 @@
 package bbgo
 
 import (
+	"encoding/json"
 	"github.com/c9s/bbgo/pkg/bbgo/types"
 	"github.com/stretchr/testify/assert"
+	"io/ioutil"
 	"testing"
 )
 
 func TestStockManager(t *testing.T) {
 
+	t.Run("testdata", func(t *testing.T) {
+		var trades []types.Trade
+
+		out, err := ioutil.ReadFile("testdata/btcusdt-trades.json")
+		assert.NoError(t, err)
+
+		err = json.Unmarshal(out, &trades)
+		assert.NoError(t, err)
+
+		var stockManager = &StockManager{
+			TradingFeeCurrency: "BNB",
+			Symbol:             "BTCUSDT",
+		}
+
+		_, err = stockManager.LoadTrades(trades)
+		assert.NoError(t, err)
+		assert.Equal(t, 0.72970242, stockManager.Stocks.Quantity())
+		assert.NotEmpty(t, stockManager.Stocks)
+		assert.Equal(t, 0, len(stockManager.PendingSells))
+
+	})
+
 	t.Run("stock", func(t *testing.T) {
 		var trades = []types.Trade{
 			{Symbol: "BTCUSDT", Price: 9100.0, Quantity: 0.05, IsBuyer: true},
@@ -23,6 +47,20 @@ func TestStockManager(t *testing.T) {
 		_, err := stockManager.LoadTrades(trades)
 		assert.NoError(t, err)
 		assert.Len(t, stockManager.Stocks, 2)
+		assert.Equal(t, StockSlice{
+			{
+				Symbol: "BTCUSDT",
+				Price: 9100.0,
+				Quantity: 0.05,
+				IsBuyer: true,
+			},
+			{
+				Symbol: "BTCUSDT",
+				Price: 9100.0,
+				Quantity: 0.04,
+				IsBuyer: true,
+			},
+		}, stockManager.Stocks)
 		assert.Len(t, stockManager.PendingSells, 0)
 	})
 
@@ -63,7 +101,6 @@ func TestStockManager(t *testing.T) {
 		assert.Len(t, stockManager.PendingSells, 1)
 	})
 
-
 	t.Run("loss sell", func(t *testing.T) {
 		var trades = []types.Trade{
 			{Symbol: "BTCUSDT", Price: 9100.0, Quantity: 0.05, IsBuyer: true},
@@ -79,12 +116,20 @@ func TestStockManager(t *testing.T) {
 		_, err := stockManager.LoadTrades(trades)
 		assert.NoError(t, err)
 		assert.Len(t, stockManager.Stocks, 1)
+		assert.Equal(t, StockSlice{
+			{
+				Symbol: "BTCUSDT",
+				Price: 9100.0,
+				Quantity: 0.02,
+				IsBuyer: true,
+			},
+		}, stockManager.Stocks)
 		assert.Len(t, stockManager.PendingSells, 0)
 	})
 
-	t.Run("pending sell", func(t *testing.T) {
+	t.Run("pending sell 1", func(t *testing.T) {
 		var trades = []types.Trade{
-			{Symbol: "BTCUSDT", Price: 9200.0, Quantity: 0.02, IsBuyer: false},
+			{Symbol: "BTCUSDT", Price: 9200.0, Quantity: 0.02},
 			{Symbol: "BTCUSDT", Price: 9100.0, Quantity: 0.05, IsBuyer: true},
 		}
 
@@ -96,8 +141,40 @@ func TestStockManager(t *testing.T) {
 		_, err := stockManager.LoadTrades(trades)
 		assert.NoError(t, err)
 		assert.Len(t, stockManager.Stocks, 1)
+		assert.Equal(t, StockSlice{
+			{
+				Symbol: "BTCUSDT",
+				Price: 9100.0,
+				Quantity: 0.03,
+				IsBuyer: true,
+			},
+		}, stockManager.Stocks)
 		assert.Len(t, stockManager.PendingSells, 0)
 	})
 
+	t.Run("pending sell 2", func(t *testing.T) {
+		var trades = []types.Trade{
+			{Symbol: "BTCUSDT", Price: 9200.0, Quantity: 0.1},
+			{Symbol: "BTCUSDT", Price: 9100.0, Quantity: 0.05, IsBuyer: true},
+		}
+
+		var stockManager = &StockManager{
+			TradingFeeCurrency: "BNB",
+			Symbol:             "BTCUSDT",
+		}
+
+		_, err := stockManager.LoadTrades(trades)
+		assert.NoError(t, err)
+		assert.Len(t, stockManager.Stocks, 0)
+		assert.Len(t, stockManager.PendingSells, 1)
+		assert.Equal(t, StockSlice{
+			{
+				Symbol: "BTCUSDT",
+				Price: 9200.0,
+				Quantity: 0.05,
+				IsBuyer: false,
+			},
+		}, stockManager.PendingSells)
+	})
 
 }