interact: fix interact tests for session

This commit is contained in:
c9s 2022-01-16 00:50:43 +08:00
parent b49fc182dc
commit b80f481e7d
3 changed files with 42 additions and 37 deletions

View File

@ -724,8 +724,7 @@ func (environ *Environment) setupTelegram(userConfig *Config, telegramBotToken s
} else {
for _, session := range sessions {
if session.IsAuthorized() {
notifier.OwnerChat = session.Chat
notifier.Owner = session.User
notifier.AddChat(session.Chat)
}
}
@ -734,6 +733,10 @@ func (environ *Environment) setupTelegram(userConfig *Config, telegramBotToken s
}
messenger.OnAuthorized(func(userSession *interact.TelegramSession) {
if userSession.IsAuthorized() {
notifier.AddChat(userSession.Chat)
}
log.Infof("user session %d got authorized, saving telegram sessions...", userSession.User.ID)
if err := sessionStore.Save(messenger.Sessions()); err != nil {
log.WithError(err).Errorf("telegram session save error")

View File

@ -113,23 +113,27 @@ func TestCustomInteraction(t *testing.T) {
err = globalInteraction.init()
assert.NoError(t, err)
m := &tb.Message{}
err = globalInteraction.runCommand("/closePosition", []string{}, telegram.newReply(m))
m := &tb.Message{
Chat: &tb.Chat{ID: 22},
Sender: &tb.User{ID: 999},
}
session := telegram.loadSession(m)
err = globalInteraction.runCommand(session, "/closePosition", []string{}, telegram.newReply(session))
assert.NoError(t, err)
assert.Equal(t, State("/closePosition_1"), globalInteraction.currentState)
assert.Equal(t, State("/closePosition_1"), session.CurrentState)
err = globalInteraction.handleResponse("BTCUSDT", telegram.newReply(m))
err = globalInteraction.handleResponse(session, "BTCUSDT", telegram.newReply(session))
assert.NoError(t, err)
assert.Equal(t, State("/closePosition_2"), globalInteraction.currentState)
assert.Equal(t, State("/closePosition_2"), session.CurrentState)
err = globalInteraction.handleResponse("0.20", telegram.newReply(m))
err = globalInteraction.handleResponse(session, "0.20", telegram.newReply(session))
assert.NoError(t, err)
assert.Equal(t, State("/closePosition_3"), globalInteraction.currentState)
assert.Equal(t, State("/closePosition_3"), session.CurrentState)
err = globalInteraction.handleResponse("true", telegram.newReply(m))
err = globalInteraction.handleResponse(session, "true", telegram.newReply(session))
assert.NoError(t, err)
assert.Equal(t, State("public"), globalInteraction.currentState)
assert.Equal(t, State("public"), session.CurrentState)
assert.Equal(t, closePositionTask{
symbol: "BTCUSDT",

View File

@ -16,13 +16,11 @@ var log = logrus.WithField("service", "telegram")
type Notifier struct {
bot *telebot.Bot
// Subscribers stores the Chat objects for broadcasting
Subscribers map[int64]time.Time `json:"chats"`
// Subscribers stores the Chat objects for broadcasting public notification
Subscribers map[int64]time.Time `json:"subscribers"`
// Owner
// when owner and owner chat is not nil, notification will send to the owner chat
Owner *telebot.User `json:"owner"`
OwnerChat *telebot.Chat `json:"chat"`
// Chats are the private chats that we will send private notification
Chats map[int64]*telebot.Chat `json:"chats"`
broadcast bool
}
@ -37,7 +35,11 @@ func UseBroadcast() Option {
// New
func New(bot *telebot.Bot, options ...Option) *Notifier {
notifier := &Notifier{ bot: bot }
notifier := &Notifier{
bot: bot,
Chats: make(map[int64]*telebot.Chat),
Subscribers: make(map[int64]time.Time),
}
for _, o := range options {
o(notifier)
@ -102,14 +104,25 @@ func (n *Notifier) NotifyTo(channel string, obj interface{}, args ...interface{}
for _, text := range texts {
n.Broadcast(text)
}
} else if n.OwnerChat != nil {
n.SendToOwner(message)
for _, text := range texts {
n.SendToOwner(text)
} else if n.Chats != nil {
for _, chat := range n.Chats {
if _, err := n.bot.Send(chat, message); err != nil {
log.WithError(err).Error("telegram send error")
}
for _, text := range texts {
if _, err := n.bot.Send(chat, text); err != nil {
log.WithError(err).Error("telegram send error")
}
}
}
}
}
func (n *Notifier) AddChat(c *telebot.Chat) {
n.Chats[c.ID] = c
}
func (n *Notifier) AddSubscriber(m *telebot.Message) {
if n.Subscribers == nil {
n.Subscribers = make(map[int64]time.Time)
@ -118,21 +131,6 @@ func (n *Notifier) AddSubscriber(m *telebot.Message) {
n.Subscribers[m.Chat.ID] = m.Time()
}
func (n *Notifier) SetOwner(owner *telebot.User, chat *telebot.Chat) {
n.Owner = owner
n.OwnerChat = chat
}
func (n *Notifier) SendToOwner(message string) {
if n.OwnerChat == nil {
return
}
if _, err := n.bot.Send(n.OwnerChat, message); err != nil {
log.WithError(err).Error("telegram send error")
}
}
func (n *Notifier) Broadcast(message string) {
if n.Subscribers == nil {
return