bbgo_origin/pkg/exchange/batch/trade_test.go

package batch

import (
	"context"
	"sync"
	"testing"
	"time"

	"github.com/pkg/errors"
	"github.com/stretchr/testify/assert"
	"go.uber.org/mock/gomock"

	"github.com/c9s/bbgo/pkg/types"
	"github.com/c9s/bbgo/pkg/types/mocks"
)

func Test_TradeBatchQuery(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	var (
		ctx          = context.Background()
		timeNow      = time.Now()
		startTime    = timeNow.Add(-24 * time.Hour)
		endTime      = timeNow
		expSymbol    = "BTCUSDT"
		queryTrades1 = []types.Trade{
			{
				ID:   1,
				Time: types.Time(startTime.Add(time.Hour)),
			},
		}
		queryTrades2 = []types.Trade{
			{
				ID:   2,
				Time: types.Time(startTime.Add(time.Hour)),
			},
			{
				ID:   3,
				Time: types.Time(startTime.Add(2 * time.Hour)),
			},
		}
		emptyTrades = []types.Trade{}
		allRes      = append(queryTrades1, queryTrades2...)
	)

	t.Run("succeeds", func(t *testing.T) {
		var (
			expOptions = &types.TradeQueryOptions{
				StartTime:   &startTime,
				EndTime:     &endTime,
				LastTradeID: 0,
				Limit:       50,
			}
			mockExchange = mocks.NewMockExchangeTradeHistoryService(ctrl)
		)

		mockExchange.EXPECT().QueryTrades(ctx, expSymbol, expOptions).DoAndReturn(
			func(ctx context.Context, symbol string, options *types.TradeQueryOptions) ([]types.Trade, error) {
				assert.Equal(t, startTime, *options.StartTime)
				assert.Equal(t, endTime, *options.EndTime)
				assert.Equal(t, uint64(0), options.LastTradeID)
				assert.Equal(t, expOptions.Limit, options.Limit)
				return queryTrades1, nil
			}).Times(1)
		mockExchange.EXPECT().QueryTrades(ctx, expSymbol, expOptions).DoAndReturn(
			func(ctx context.Context, symbol string, options *types.TradeQueryOptions) ([]types.Trade, error) {
				assert.Equal(t, queryTrades1[0].Time.Time(), *options.StartTime)
				assert.Equal(t, endTime, *options.EndTime)
				assert.Equal(t, queryTrades1[0].ID, options.LastTradeID)
				assert.Equal(t, expOptions.Limit, options.Limit)
				return queryTrades2, nil
			}).Times(1)
		mockExchange.EXPECT().QueryTrades(ctx, expSymbol, expOptions).DoAndReturn(
			func(ctx context.Context, symbol string, options *types.TradeQueryOptions) ([]types.Trade, error) {
				assert.Equal(t, queryTrades2[1].Time.Time(), *options.StartTime)
				assert.Equal(t, endTime, *options.EndTime)
				assert.Equal(t, queryTrades2[1].ID, options.LastTradeID)
				assert.Equal(t, expOptions.Limit, options.Limit)
				return emptyTrades, nil
			}).Times(1)

		tradeBatchQuery := &TradeBatchQuery{ExchangeTradeHistoryService: mockExchange}

		resCh, errC := tradeBatchQuery.Query(ctx, expSymbol, expOptions)
		wg := sync.WaitGroup{}
		wg.Add(1)
		rcvCount := 0
		go func() {
			defer wg.Done()
			for ch := range resCh {
				assert.Equal(t, allRes[rcvCount], ch)
				rcvCount++
			}
			assert.NoError(t, <-errC)
		}()
		wg.Wait()
		assert.Equal(t, rcvCount, len(allRes))
	})

	t.Run("failed to call query trades", func(t *testing.T) {
		var (
			expOptions = &types.TradeQueryOptions{
				StartTime:   &startTime,
				EndTime:     &endTime,
				LastTradeID: 0,
				Limit:       50,
			}
			mockExchange = mocks.NewMockExchangeTradeHistoryService(ctrl)
			unknownErr   = errors.New("unknown err")
		)

		mockExchange.EXPECT().QueryTrades(ctx, expSymbol, expOptions).DoAndReturn(
			func(ctx context.Context, symbol string, options *types.TradeQueryOptions) ([]types.Trade, error) {
				assert.Equal(t, startTime, *options.StartTime)
				assert.Equal(t, endTime, *options.EndTime)
				assert.Equal(t, uint64(0), options.LastTradeID)
				assert.Equal(t, expOptions.Limit, options.Limit)
				return nil, unknownErr
			}).Times(1)

		tradeBatchQuery := &TradeBatchQuery{ExchangeTradeHistoryService: mockExchange}

		resCh, errC := tradeBatchQuery.Query(ctx, expSymbol, expOptions)
		wg := sync.WaitGroup{}
		wg.Add(1)
		rcvCount := 0
		go func() {
			defer wg.Done()
			for ch := range resCh {
				assert.Equal(t, allRes[rcvCount], ch)
				rcvCount++
			}
			assert.Equal(t, 0, rcvCount)
			assert.Equal(t, unknownErr, <-errC)
		}()
		wg.Wait()
		assert.Equal(t, rcvCount, 0)
	})
}