Merge pull request #1274 from bailantaotao/edwin/add-account-info

FEATURE: [bybit] add balance snapshot event

pkg/exchange/bybit/stream.go

@@ -32,7 +32,8 @@ type Stream struct {
 	key, secret string
 	types.StandardStream
 
-	bookEventCallbacks []func(e BookEvent)
+	bookEventCallbacks   []func(e BookEvent)
+	walletEventCallbacks []func(e []*WalletEvent)
 }
 
 func NewStream(key, secret string) *Stream {
@@ -50,6 +51,7 @@ func NewStream(key, secret string) *Stream {
 
 	stream.OnConnect(stream.handlerConnect)
 	stream.OnBookEvent(stream.handleBookEvent)
+	stream.OnWalletEvent(stream.handleWalletEvent)
 	return stream
 }
 
@@ -72,6 +74,9 @@ func (s *Stream) dispatchEvent(event interface{}) {
 
 	case *BookEvent:
 		s.EmitBookEvent(*e)
+
+	case []*WalletEvent:
+		s.EmitWalletEvent(e)
 	}
 }
 
@@ -98,6 +103,10 @@ func (s *Stream) parseWebSocketEvent(in []byte) (interface{}, error) {
 
 			book.Type = e.WebSocketTopicEvent.Type
 			return &book, nil
+
+		case TopicTypeWallet:
+			var wallets []*WalletEvent
+			return wallets, json.Unmarshal(e.WebSocketTopicEvent.Data, &wallets)
 		}
 	}
 
@@ -164,6 +173,7 @@ func (s *Stream) handlerConnect() {
 			Args: topics,
 		}); err != nil {
 			log.WithError(err).Error("failed to send subscription request")
+			return
 		}
 	} else {
 		expires := strconv.FormatInt(time.Now().Add(wsAuthRequest).In(time.UTC).UnixMilli(), 10)
@@ -177,6 +187,17 @@ func (s *Stream) handlerConnect() {
 			},
 		}); err != nil {
 			log.WithError(err).Error("failed to auth request")
+			return
+		}
+
+		if err := s.Conn.WriteJSON(WebsocketOp{
+			Op: WsOpTypeSubscribe,
+			Args: []string{
+				string(TopicTypeWallet),
+			},
+		}); err != nil {
+			log.WithError(err).Error("failed to send subscription request")
+			return
 		}
 	}
 }
@@ -206,3 +227,22 @@ func (s *Stream) handleBookEvent(e BookEvent) {
 		s.EmitBookUpdate(orderBook)
 	}
 }
+
+func (s *Stream) handleWalletEvent(events []*WalletEvent) {
+	bm := types.BalanceMap{}
+	for _, event := range events {
+		if event.AccountType != AccountTypeSpot {
+			return
+		}
+
+		for _, obj := range event.Coins {
+			bm[obj.Coin] = types.Balance{
+				Currency:  obj.Coin,
+				Available: obj.Free,
+				Locked:    obj.Locked,
+			}
+		}
+	}
+
+	s.StandardStream.EmitBalanceSnapshot(bm)
+}

pkg/exchange/bybit/stream_callbacks.go

@@ -13,3 +13,13 @@ func (s *Stream) EmitBookEvent(e BookEvent) {
 		cb(e)
 	}
 }
+
+func (s *Stream) OnWalletEvent(cb func(e []*WalletEvent)) {
+	s.walletEventCallbacks = append(s.walletEventCallbacks, cb)
+}
+
+func (s *Stream) EmitWalletEvent(e []*WalletEvent) {
+	for _, cb := range s.walletEventCallbacks {
+		cb(e)
+	}
+}

pkg/exchange/bybit/stream_test.go

@@ -36,6 +36,35 @@ func TestStream(t *testing.T) {
 		c := make(chan struct{})
 		<-c
 	})
+
+	t.Run("book test", func(t *testing.T) {
+		s.Subscribe(types.BookChannel, "BTCUSDT", types.SubscribeOptions{
+			Depth: types.DepthLevel50,
+		})
+		s.SetPublicOnly()
+		err := s.Connect(context.Background())
+		assert.NoError(t, err)
+
+		s.OnBookSnapshot(func(book types.SliceOrderBook) {
+			t.Log("got snapshot", book)
+		})
+		s.OnBookUpdate(func(book types.SliceOrderBook) {
+			t.Log("got update", book)
+		})
+		c := make(chan struct{})
+		<-c
+	})
+
+	t.Run("wallet test", func(t *testing.T) {
+		err := s.Connect(context.Background())
+		assert.NoError(t, err)
+
+		s.OnBalanceSnapshot(func(balances types.BalanceMap) {
+			t.Log("got snapshot", balances)
+		})
+		c := make(chan struct{})
+		<-c
+	})
 }
 
 func TestStream_parseWebSocketEvent(t *testing.T) {

pkg/exchange/bybit/types.go

@@ -64,7 +64,9 @@ func (w *WebSocketOpEvent) IsValid() error {
 		}
 		return nil
 	case WsOpTypeSubscribe:
-		if !w.Success || WsOpType(w.RetMsg) != WsOpTypeSubscribe {
+		// in the public channel, you can get RetMsg = 'subscribe', but in the private channel, you cannot.
+		// so, we only verify that success is true.
+		if !w.Success {
 			return fmt.Errorf("unexpected response result: %+v", w)
 		}
 		return nil
@@ -77,6 +79,7 @@ type TopicType string
 
 const (
 	TopicTypeOrderBook TopicType = "orderbook"
+	TopicTypeWallet    TopicType = "wallet"
 )
 
 type DataType string
@@ -137,3 +140,64 @@ func getTopicType(topic string) TopicType {
 	}
 	return TopicType(slice[0])
 }
+
+type AccountType string
+
+const AccountTypeSpot AccountType = "SPOT"
+
+type WalletEvent struct {
+	AccountType            AccountType      `json:"accountType"`
+	AccountIMRate          fixedpoint.Value `json:"accountIMRate"`
+	AccountMMRate          fixedpoint.Value `json:"accountMMRate"`
+	TotalEquity            fixedpoint.Value `json:"totalEquity"`
+	TotalWalletBalance     fixedpoint.Value `json:"totalWalletBalance"`
+	TotalMarginBalance     fixedpoint.Value `json:"totalMarginBalance"`
+	TotalAvailableBalance  fixedpoint.Value `json:"totalAvailableBalance"`
+	TotalPerpUPL           fixedpoint.Value `json:"totalPerpUPL"`
+	TotalInitialMargin     fixedpoint.Value `json:"totalInitialMargin"`
+	TotalMaintenanceMargin fixedpoint.Value `json:"totalMaintenanceMargin"`
+	// Account LTV: account total borrowed size / (account total equity + account total borrowed size).
+	// In non-unified mode & unified (inverse) & unified (isolated_margin), the field will be returned as an empty string.
+	AccountLTV fixedpoint.Value `json:"accountLTV"`
+	Coins      []struct {
+		Coin string `json:"coin"`
+		// Equity of current coin
+		Equity fixedpoint.Value `json:"equity"`
+		// UsdValue of current coin. If this coin cannot be collateral, then it is 0
+		UsdValue fixedpoint.Value `json:"usdValue"`
+		// WalletBalance of current coin
+		WalletBalance fixedpoint.Value `json:"walletBalance"`
+		// Free available balance for Spot wallet. This is a unique field for Normal SPOT
+		Free fixedpoint.Value
+		// Locked balance for Spot wallet. This is a unique field for Normal SPOT
+		Locked fixedpoint.Value
+		// Available amount to withdraw of current coin
+		AvailableToWithdraw fixedpoint.Value `json:"availableToWithdraw"`
+		// Available amount to borrow of current coin
+		AvailableToBorrow fixedpoint.Value `json:"availableToBorrow"`
+		// Borrow amount of current coin
+		BorrowAmount fixedpoint.Value `json:"borrowAmount"`
+		// Accrued interest
+		AccruedInterest fixedpoint.Value `json:"accruedInterest"`
+		// Pre-occupied margin for order. For portfolio margin mode, it returns ""
+		TotalOrderIM fixedpoint.Value `json:"totalOrderIM"`
+		// Sum of initial margin of all positions + Pre-occupied liquidation fee. For portfolio margin mode, it returns ""
+		TotalPositionIM fixedpoint.Value `json:"totalPositionIM"`
+		// Sum of maintenance margin for all positions. For portfolio margin mode, it returns ""
+		TotalPositionMM fixedpoint.Value `json:"totalPositionMM"`
+		// Unrealised P&L
+		UnrealisedPnl fixedpoint.Value `json:"unrealisedPnl"`
+		// Cumulative Realised P&L
+		CumRealisedPnl fixedpoint.Value `json:"cumRealisedPnl"`
+		// Bonus. This is a unique field for UNIFIED account
+		Bonus fixedpoint.Value `json:"bonus"`
+		// Whether it can be used as a margin collateral currency (platform)
+		// - When marginCollateral=false, then collateralSwitch is meaningless
+		// -  This is a unique field for UNIFIED account
+		CollateralSwitch bool `json:"collateralSwitch"`
+		// Whether the collateral is turned on by user (user)
+		// - When marginCollateral=true, then collateralSwitch is meaningful
+		// - This is a unique field for UNIFIED account
+		MarginCollateral bool `json:"marginCollateral"`
+	} `json:"coin"`
+}

pkg/exchange/bybit/types_test.go

@@ -173,7 +173,7 @@ func Test_WebSocketEventIsValid(t *testing.T) {
 		assert.Equal(t, fmt.Errorf("unexpected op type: %+v", w), w.IsValid())
 	})
 
-	t.Run("[subscribe] valid", func(t *testing.T) {
+	t.Run("[subscribe] valid with public channel", func(t *testing.T) {
 		expRetMsg := "subscribe"
 		w := &WebSocketOpEvent{
 			Success: true,
@@ -186,6 +186,18 @@ func Test_WebSocketEventIsValid(t *testing.T) {
 		assert.NoError(t, w.IsValid())
 	})
 
+	t.Run("[subscribe] valid with private channel", func(t *testing.T) {
+		w := &WebSocketOpEvent{
+			Success: true,
+			RetMsg:  "",
+			ReqId:   "",
+			ConnId:  "test-conndid",
+			Op:      WsOpTypeSubscribe,
+			Args:    nil,
+		}
+		assert.NoError(t, w.IsValid())
+	})
+
 	t.Run("[subscribe] un-succeeds", func(t *testing.T) {
 		expRetMsg := ""
 		w := &WebSocketOpEvent{