Merge pull request #2029 from anywhy/improve_risk_service

Author: c9s

pkg/exchange/binance/exchange_test.go

@@ -6,6 +6,7 @@ import (
 	"strings"
 	"testing"
 
+	"github.com/c9s/bbgo/pkg/testutil"
 	"github.com/stretchr/testify/assert"
 )
 
@@ -32,3 +33,40 @@ func Test_new(t *testing.T) {
 		assert.NoError(t, err)
 	}
 }
+
+func Test_QueryPositionRisk(t *testing.T) {
+	key, secret, ok := testutil.IntegrationTestConfigured(t, "BINANCE")
+	if !ok {
+		t.SkipNow()
+		return
+	}
+
+	ex := New(key, secret)
+	ex.UseFutures()
+
+	assert.NotEmpty(t, ex)
+	ctx := context.Background()
+
+	// Test case 1: Query all positions
+	positions, err := ex.QueryPositionRisk(ctx)
+	assert.NoError(t, err)
+
+	// Test case 2: Query specific symbol position
+	positions, err = ex.QueryPositionRisk(ctx, "BTCUSDT")
+	assert.NoError(t, err)
+	if len(positions) > 0 {
+		assert.Equal(t, "BTCUSDT", positions[0].Symbol)
+	}
+
+	// Test case 3: Query multiple symbols
+	positions, err = ex.QueryPositionRisk(ctx, "BTCUSDT", "ETHUSDT")
+	assert.NoError(t, err)
+	if len(positions) > 0 {
+		symbols := make(map[string]bool)
+		for _, pos := range positions {
+			symbols[pos.Symbol] = true
+		}
+		assert.True(t, symbols["BTCUSDT"] || symbols["ETHUSDT"])
+	}
+
+}

pkg/exchange/binance/futures.go

@@ -412,23 +412,43 @@ func (e *Exchange) SetLeverage(ctx context.Context, symbol string, leverage int)
 	return fmt.Errorf("not supported set leverage")
 }
 
-func (e *Exchange) QueryPositionRisk(ctx context.Context, symbol string) ([]types.PositionRisk, error) {
-	var (
-		positions []binanceapi.FuturesPositionRisk
-		err       error
-	)
+func (e *Exchange) QueryPositionRisk(ctx context.Context, symbol ...string) ([]types.PositionRisk, error) {
+	if !e.IsFutures {
+		return nil, fmt.Errorf("not supported for non-futures exchange")
+	}
 
-	if e.IsFutures {
-		req := e.futuresClient2.NewFuturesGetPositionRisksRequest()
-		if len(symbol) > 0 {
-			req.Symbol(symbol)
-		}
-		if positions, err = req.Do(ctx); err != nil {
+	req := e.futuresClient2.NewFuturesGetPositionRisksRequest()
+	if len(symbol) == 1 {
+		req.Symbol(symbol[0])
+		positions, err := req.Do(ctx)
+		if err != nil {
 			return nil, err
 		}
+		return toGlobalPositionRisk(positions), nil
+	}
+
+	positions, err := req.Do(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	if len(symbol) == 0 {
+		return toGlobalPositionRisk(positions), nil
+	}
+
+	symbolSet := make(map[string]struct{}, len(symbol))
+	for _, s := range symbol {
+		symbolSet[s] = struct{}{}
+	}
+
+	filteredPositions := make([]binanceapi.FuturesPositionRisk, 0, len(symbol))
+	for _, pos := range positions {
+		if _, ok := symbolSet[pos.Symbol]; ok {
+			filteredPositions = append(filteredPositions, pos)
+		}
 	}
 
-	return toGlobalPositionRisk(positions), nil
+	return toGlobalPositionRisk(filteredPositions), nil
 }
 
 func setDualSidePosition(req *futures.CreateOrderService, order types.SubmitOrder) {

pkg/exchange/okex/exchange_test.go

@@ -693,3 +693,68 @@ func TestExchange_QueryMarkets_SymbolCache(t *testing.T) {
 		assert.NotContains(t, spotSymbolMap, "ETHUSDT")
 	})
 }
+
+func TestExchange_QueryPositionRisk(t *testing.T) {
+	key, secret, passphrase, ok := testutil.IntegrationTestWithPassphraseConfigured(t, "OKEX")
+	if !ok {
+		t.SkipNow()
+		return
+	}
+
+	ex := New(key, secret, passphrase)
+	ex.UseFutures()
+
+	ctx := context.Background()
+
+	t.Run("query all positions", func(t *testing.T) {
+		// Test querying all positions
+		positions, err := ex.QueryPositionRisk(ctx)
+		assert.NoError(t, err)
+		assert.NotNil(t, positions)
+
+		// Verify position data structure
+		for _, pos := range positions {
+			assert.NotEmpty(t, pos.Symbol)
+			assert.NotEmpty(t, pos.EntryPrice)
+			assert.NotEmpty(t, pos.PositionAmount)
+			assert.NotEmpty(t, pos.Leverage)
+			assert.NotEmpty(t, pos.MarkPrice)
+			assert.NotEmpty(t, pos.LiquidationPrice)
+		}
+	})
+
+	t.Run("query specific symbol position", func(t *testing.T) {
+		// Test querying specific symbol position
+		symbol := "BTCUSDT"
+		positions, err := ex.QueryPositionRisk(ctx, symbol)
+		assert.NoError(t, err)
+		assert.NotNil(t, positions)
+
+		// If there are positions, verify the symbol matches
+		if len(positions) > 0 {
+			for _, pos := range positions {
+				assert.Equal(t, symbol, pos.Symbol)
+			}
+		}
+	})
+
+	t.Run("query multiple symbols", func(t *testing.T) {
+		// Test querying multiple symbols
+		symbols := []string{"BTCUSDT", "ETHUSDT"}
+		positions, err := ex.QueryPositionRisk(ctx, symbols...)
+		assert.NoError(t, err)
+		assert.NotNil(t, positions)
+
+		// If there are positions, verify they belong to requested symbols
+		if len(positions) > 0 {
+			symbolSet := make(map[string]bool)
+			for _, symbol := range symbols {
+				symbolSet[symbol] = true
+			}
+
+			for _, pos := range positions {
+				assert.True(t, symbolSet[pos.Symbol], "Position symbol should be one of the requested symbols")
+			}
+		}
+	})
+}

pkg/exchange/okex/futures.go

@@ -3,6 +3,7 @@ package okex
 import (
 	"context"
 	"fmt"
+	"strings"
 
 	"github.com/c9s/bbgo/pkg/exchange/okex/okexapi"
 	"github.com/c9s/bbgo/pkg/types"
@@ -51,7 +52,7 @@ func (e *Exchange) SetLeverage(ctx context.Context, symbol string, leverage int)
 	return fmt.Errorf("not supported set leverage")
 }
 
-func (e *Exchange) QueryPositionRisk(ctx context.Context, symbol string) ([]types.PositionRisk, error) {
+func (e *Exchange) QueryPositionRisk(ctx context.Context, symbol ...string) ([]types.PositionRisk, error) {
 	var (
 		positions []okexapi.Position
 		err       error
@@ -61,7 +62,11 @@ func (e *Exchange) QueryPositionRisk(ctx context.Context, symbol string) ([]type
 		req := e.client.NewGetAccountPositionsRequest().
 			InstType(okexapi.InstrumentTypeSwap)
 		if len(symbol) > 0 {
-			req.InstId(toLocalSymbol(symbol, okexapi.InstrumentTypeSwap))
+			symbols := make([]string, len(symbol))
+			for i, s := range symbol {
+				symbols[i] = toLocalSymbol(s, okexapi.InstrumentTypeSwap)
+			}
+			req.InstId(strings.Join(symbols, ","))
 		}
 		if positions, err = req.Do(ctx); err != nil {
 			return nil, err

pkg/types/exchange.go

@@ -179,7 +179,7 @@ type ExchangeRewardService interface {
 
 type ExchangeRiskService interface {
 	SetLeverage(ctx context.Context, symbol string, leverage int) error
-	QueryPositionRisk(ctx context.Context, symbol string) ([]PositionRisk, error)
+	QueryPositionRisk(ctx context.Context, symbol ...string) ([]PositionRisk, error)
 }
 
 type TradeQueryOptions struct {