diff --git a/bot.go b/bot.go index 9297d16..0ab7d44 100644 --- a/bot.go +++ b/bot.go @@ -44,11 +44,11 @@ type Bot struct { errorsHandler ErrorsHandler debugHandler DebugHandler - middlewaresMx *sync.RWMutex + middlewaresMx sync.RWMutex middlewares []Middleware - handlersMx *sync.RWMutex - handlers map[string]handler + handlersMx sync.RWMutex + handlers []handler client HttpClient lastUpdateID int64 @@ -67,13 +67,9 @@ func New(token string, options ...Option) (*Bot, error) { } b := &Bot{ - url: "https://api.telegram.org", - token: token, - pollTimeout: defaultPollTimeout, - middlewaresMx: &sync.RWMutex{}, - middlewares: []Middleware{}, - handlersMx: &sync.RWMutex{}, - handlers: map[string]handler{}, + url: "https://api.telegram.org", + token: token, + pollTimeout: defaultPollTimeout, client: &http.Client{ Timeout: defaultPollTimeout, }, @@ -105,11 +101,11 @@ func New(token string, options ...Option) (*Bot, error) { // StartWebhook starts the Bot with webhook mode func (b *Bot) StartWebhook(ctx context.Context) { - wg := &sync.WaitGroup{} + wg := sync.WaitGroup{} wg.Add(b.workers) for i := 0; i < b.workers; i++ { - go b.waitUpdates(ctx, wg) + go b.waitUpdates(ctx, &wg) } wg.Wait() @@ -117,14 +113,14 @@ func (b *Bot) StartWebhook(ctx context.Context) { // Start the bot func (b *Bot) Start(ctx context.Context) { - wg := &sync.WaitGroup{} + wg := sync.WaitGroup{} wg.Add(1) - go b.getUpdates(ctx, wg) + go b.getUpdates(ctx, &wg) wg.Add(b.workers) for i := 0; i < b.workers; i++ { - go b.waitUpdates(ctx, wg) + go b.waitUpdates(ctx, &wg) } wg.Wait() diff --git a/handlers.go b/handlers.go index 2d356c0..eb4fefa 100644 --- a/handlers.go +++ b/handlers.go @@ -27,6 +27,7 @@ const ( ) type handler struct { + id string handlerType HandlerType matchType MatchType handler HandlerFunc @@ -80,12 +81,13 @@ func (b *Bot) RegisterHandlerMatchFunc(matchFunc MatchFunc, f HandlerFunc, m ... id := RandomString(16) h := handler{ + id: id, matchType: matchTypeFunc, matchFunc: matchFunc, handler: applyMiddlewares(f, m...), } - b.handlers[id] = h + b.handlers = append(b.handlers, h) return id } @@ -97,13 +99,14 @@ func (b *Bot) RegisterHandlerRegexp(handlerType HandlerType, re *regexp.Regexp, id := RandomString(16) h := handler{ + id: id, handlerType: handlerType, matchType: matchTypeRegexp, re: re, handler: applyMiddlewares(f, m...), } - b.handlers[id] = h + b.handlers = append(b.handlers, h) return id } @@ -115,13 +118,14 @@ func (b *Bot) RegisterHandler(handlerType HandlerType, pattern string, matchType id := RandomString(16) h := handler{ + id: id, handlerType: handlerType, matchType: matchType, pattern: pattern, handler: applyMiddlewares(f, m...), } - b.handlers[id] = h + b.handlers = append(b.handlers, h) return id } @@ -130,5 +134,10 @@ func (b *Bot) UnregisterHandler(id string) { b.handlersMx.Lock() defer b.handlersMx.Unlock() - delete(b.handlers, id) + for i, h := range b.handlers { + if h.id == id { + b.handlers = append(b.handlers[:i], b.handlers[i+1:]...) + return + } + } } diff --git a/handlers_test.go b/handlers_test.go index a054140..804e8ac 100644 --- a/handlers_test.go +++ b/handlers_test.go @@ -2,18 +2,27 @@ package bot import ( "regexp" - "sync" "testing" "github.com/go-telegram/bot/models" ) -func Test_match_func(t *testing.T) { - b := &Bot{ - handlersMx: &sync.RWMutex{}, - handlers: map[string]handler{}, +func findHandler(b *Bot, id string) *handler { + b.handlersMx.RLock() + defer b.handlersMx.RUnlock() + + for _, h := range b.handlers { + if h.id == id { + return &h + } } + return nil +} + +func Test_match_func(t *testing.T) { + b := &Bot{} + var called bool id := b.RegisterHandlerMatchFunc(func(update *models.Update) bool { @@ -24,7 +33,7 @@ func Test_match_func(t *testing.T) { return true }, nil) - h := b.handlers[id] + h := findHandler(b, id) res := h.match(&models.Update{ID: 42}) if !called { @@ -36,14 +45,11 @@ func Test_match_func(t *testing.T) { } func Test_match_exact(t *testing.T) { - b := &Bot{ - handlersMx: &sync.RWMutex{}, - handlers: map[string]handler{}, - } + b := &Bot{} id := b.RegisterHandler(HandlerTypeMessageText, "xxx", MatchTypeExact, nil) - h := b.handlers[id] + h := findHandler(b, id) res := h.match(&models.Update{Message: &models.Message{Text: "zzz"}}) if res { @@ -57,14 +63,11 @@ func Test_match_exact(t *testing.T) { } func Test_match_prefix(t *testing.T) { - b := &Bot{ - handlersMx: &sync.RWMutex{}, - handlers: map[string]handler{}, - } + b := &Bot{} id := b.RegisterHandler(HandlerTypeCallbackQueryData, "abc", MatchTypePrefix, nil) - h := b.handlers[id] + h := findHandler(b, id) res := h.match(&models.Update{CallbackQuery: &models.CallbackQuery{Data: "xabcdef"}}) if res { @@ -78,14 +81,11 @@ func Test_match_prefix(t *testing.T) { } func Test_match_contains(t *testing.T) { - b := &Bot{ - handlersMx: &sync.RWMutex{}, - handlers: map[string]handler{}, - } + b := &Bot{} id := b.RegisterHandler(HandlerTypeCallbackQueryData, "abc", MatchTypeContains, nil) - h := b.handlers[id] + h := findHandler(b, id) res := h.match(&models.Update{CallbackQuery: &models.CallbackQuery{Data: "xxabxx"}}) if res { @@ -99,16 +99,13 @@ func Test_match_contains(t *testing.T) { } func Test_match_regexp(t *testing.T) { - b := &Bot{ - handlersMx: &sync.RWMutex{}, - handlers: map[string]handler{}, - } + b := &Bot{} re := regexp.MustCompile("^[a-z]+") id := b.RegisterHandlerRegexp(HandlerTypeCallbackQueryData, re, nil) - h := b.handlers[id] + h := findHandler(b, id) res := h.match(&models.Update{CallbackQuery: &models.CallbackQuery{Data: "123abc"}}) if res { @@ -122,14 +119,11 @@ func Test_match_regexp(t *testing.T) { } func Test_match_invalid_type(t *testing.T) { - b := &Bot{ - handlersMx: &sync.RWMutex{}, - handlers: map[string]handler{}, - } + b := &Bot{} id := b.RegisterHandler(-1, "", -1, nil) - h := b.handlers[id] + h := findHandler(b, id) res := h.match(&models.Update{CallbackQuery: &models.CallbackQuery{Data: "123abc"}}) if res { @@ -138,10 +132,7 @@ func Test_match_invalid_type(t *testing.T) { } func TestBot_RegisterUnregisterHandler(t *testing.T) { - b := &Bot{ - handlersMx: &sync.RWMutex{}, - handlers: map[string]handler{}, - } + b := &Bot{} id1 := b.RegisterHandler(HandlerTypeCallbackQueryData, "", MatchTypeExact, nil) id2 := b.RegisterHandler(HandlerTypeCallbackQueryData, "", MatchTypeExact, nil) @@ -149,10 +140,10 @@ func TestBot_RegisterUnregisterHandler(t *testing.T) { if len(b.handlers) != 2 { t.Fatalf("unexpected handlers len") } - if _, ok := b.handlers[id1]; !ok { + if h := findHandler(b, id1); h == nil { t.Fatalf("handler not found") } - if _, ok := b.handlers[id2]; !ok { + if h := findHandler(b, id2); h == nil { t.Fatalf("handler not found") } @@ -160,23 +151,20 @@ func TestBot_RegisterUnregisterHandler(t *testing.T) { if len(b.handlers) != 1 { t.Fatalf("unexpected handlers len") } - if _, ok := b.handlers[id1]; ok { + if h := findHandler(b, id1); h != nil { t.Fatalf("handler found") } - if _, ok := b.handlers[id2]; !ok { + if h := findHandler(b, id2); h == nil { t.Fatalf("handler not found") } } func Test_match_exact_game(t *testing.T) { - b := &Bot{ - handlersMx: &sync.RWMutex{}, - handlers: map[string]handler{}, - } + b := &Bot{} id := b.RegisterHandler(HandlerTypeCallbackQueryGameShortName, "xxx", MatchTypeExact, nil) - h := b.handlers[id] + h := findHandler(b, id) u := models.Update{ ID: 42, CallbackQuery: &models.CallbackQuery{ diff --git a/process_update_test.go b/process_update_test.go index a5f4bfb..c95d7ee 100644 --- a/process_update_test.go +++ b/process_update_test.go @@ -2,7 +2,6 @@ package bot import ( "context" - "sync" "testing" "github.com/go-telegram/bot/models" @@ -38,8 +37,6 @@ func TestProcessUpdate(t *testing.T) { bot := &Bot{ defaultHandlerFunc: h, middlewares: []Middleware{}, - handlersMx: &sync.RWMutex{}, - handlers: map[string]handler{}, } ctx := context.Background() @@ -66,8 +63,6 @@ func TestProcessUpdate_WithMiddlewares(t *testing.T) { bot := &Bot{ defaultHandlerFunc: h, middlewares: []Middleware{middleware}, - handlersMx: &sync.RWMutex{}, - handlers: map[string]handler{}, } ctx := context.Background() @@ -93,8 +88,6 @@ func TestProcessUpdate_WithMatchTypeFunc(t *testing.T) { bot := &Bot{ defaultHandlerFunc: h1, - handlersMx: &sync.RWMutex{}, - handlers: map[string]handler{}, } bot.RegisterHandlerMatchFunc(m, h2) @@ -116,17 +109,16 @@ func Test_findHandler(t *testing.T) { bot := &Bot{ defaultHandlerFunc: h, - handlersMx: &sync.RWMutex{}, - handlers: map[string]handler{}, } // Register a handler - bot.handlers["test"] = handler{ + bot.handlers = append(bot.handlers, handler{ + id: "test", handlerType: HandlerTypeMessageText, matchType: MatchTypeExact, pattern: "test", handler: h, - } + }) ctx := context.Background() upd := &models.Update{Message: &models.Message{Text: "test"}} @@ -147,8 +139,6 @@ func Test_findHandler_Default(t *testing.T) { bot := &Bot{ defaultHandlerFunc: h, - handlersMx: &sync.RWMutex{}, - handlers: map[string]handler{}, } ctx := context.Background()