From d63cc42867d33260db5b31d0302888f9a61e956d Mon Sep 17 00:00:00 2001
From: c9s <yoan@maicoin.com>
Date: Fri, 7 Jan 2022 01:17:07 +0800
Subject: [PATCH] bbgo: add trade collector test

---
 pkg/bbgo/tradecollector.go      |  9 +++--
 pkg/bbgo/tradecollector_test.go | 68 +++++++++++++++++++++++++++++++++
 2 files changed, 74 insertions(+), 3 deletions(-)
 create mode 100644 pkg/bbgo/tradecollector_test.go

diff --git a/pkg/bbgo/tradecollector.go b/pkg/bbgo/tradecollector.go
index 2f078b6f5..1104d2cc2 100644
--- a/pkg/bbgo/tradecollector.go
+++ b/pkg/bbgo/tradecollector.go
@@ -50,7 +50,9 @@ func (c *TradeCollector) BindStreamForBackground(stream types.Stream) {
 }
 
 func (c *TradeCollector) BindStream(stream types.Stream) {
-	stream.OnTradeUpdate(c.ProcessTrade)
+	stream.OnTradeUpdate(func(trade types.Trade) {
+		c.ProcessTrade(trade)
+	})
 }
 
 // Emit triggers the trade processing (position update)
@@ -66,7 +68,6 @@ func (c *TradeCollector) Emit() {
 func (c *TradeCollector) Process() bool {
 	positionChanged := false
 	c.tradeStore.Filter(func(trade types.Trade) bool {
-		return c.processTrade(trade)
 		key := trade.Key()
 
 		// if it's already done, remove the trade from the trade store
@@ -112,10 +113,12 @@ func (c *TradeCollector) processTrade(trade types.Trade) bool {
 	return false
 }
 
-func (c *TradeCollector) ProcessTrade(trade types.Trade) {
+func (c *TradeCollector) ProcessTrade(trade types.Trade) bool {
 	if !c.processTrade(trade) {
 		c.tradeStore.Add(trade)
+		return false
 	}
+	return true
 }
 
 // Run is a goroutine executed in the background
diff --git a/pkg/bbgo/tradecollector_test.go b/pkg/bbgo/tradecollector_test.go
new file mode 100644
index 000000000..d4a1bd05f
--- /dev/null
+++ b/pkg/bbgo/tradecollector_test.go
@@ -0,0 +1,68 @@
+package bbgo
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+
+	"github.com/c9s/bbgo/pkg/types"
+)
+
+func TestTradeCollector(t *testing.T) {
+	symbol := "BTCUSDT"
+	position := types.NewPosition(symbol, "BTC", "USDT")
+	orderStore := NewOrderStore(symbol)
+	collector := NewTradeCollector(symbol, position, orderStore)
+	assert.NotNil(t, collector)
+
+	matched := collector.ProcessTrade(types.Trade{
+		ID:            1,
+		OrderID:       399,
+		Exchange:      types.ExchangeBinance,
+		Price:         40000.0,
+		Quantity:      1.0,
+		QuoteQuantity: 40000.0,
+		Symbol:        "BTCUSDT",
+		Side:          types.SideTypeBuy,
+		IsBuyer:       true,
+		Fee:           0,
+		FeeCurrency:   "",
+	})
+	assert.False(t, matched, "should be added to the trade store")
+	assert.Equal(t, 1, len(collector.tradeStore.Trades()), "should have one trade in the trade store")
+
+	orderStore.Add(types.Order{
+		SubmitOrder:      types.SubmitOrder{
+			Symbol:           "BTCUSDT",
+			Side:             types.SideTypeBuy,
+			Type:             types.OrderTypeLimit,
+			Quantity:         1.0,
+			Price:            40000.0,
+		},
+		Exchange:         types.ExchangeBinance,
+		OrderID:          399,
+		Status:           types.OrderStatusFilled,
+		ExecutedQuantity: 1.0,
+		IsWorking:        false,
+	})
+
+	matched = collector.Process()
+	assert.True(t, matched)
+	assert.Equal(t, 0, len(collector.tradeStore.Trades()), "the found trade should be removed from the trade store")
+
+	matched = collector.ProcessTrade(types.Trade{
+		ID:            1,
+		OrderID:       399,
+		Exchange:      types.ExchangeBinance,
+		Price:         40000.0,
+		Quantity:      1.0,
+		QuoteQuantity: 40000.0,
+		Symbol:        "BTCUSDT",
+		Side:          types.SideTypeBuy,
+		IsBuyer:       true,
+		Fee:           0,
+		FeeCurrency:   "",
+	})
+	assert.True(t, matched, "the same trade should match")
+	assert.Equal(t, 0, len(collector.tradeStore.Trades()), "the same trade should not be added to the trade store")
+}