Merge branch 'poller-fix' into develop

bot.go

@@ -189,7 +189,7 @@ func (b *Bot) Start() {
 			b.ProcessUpdate(upd)
 		// call to stop polling
 		case <-b.stop:
-			stop <- struct{}{}
+			close(stop)
 			return
 		}
 	}

bot_test.go

@@ -107,11 +107,10 @@ func TestBotStart(t *testing.T) {
 	// remove webhook to be sure that bot can poll
 	assert.NoError(t, b.RemoveWebhook())
 
-	time.AfterFunc(50*time.Millisecond, b.Stop)
-	b.Start() // stops after some delay
-	assert.Empty(t, b.stop)
+	go b.Start()
+	b.Stop()
 
-	tp := &testPoller{updates: make(chan Update, 1)}
+	tp := newTestPoller()
 	go func() {
 		tp.updates <- Update{Message: &Message{Text: "/start"}}
 	}()
@@ -123,11 +122,14 @@ func TestBotStart(t *testing.T) {
 	var ok bool
 	b.Handle("/start", func(m *Message) {
 		assert.Equal(t, m.Text, "/start")
+		tp.done <- struct{}{}
 		ok = true
 	})
 
-	time.AfterFunc(100*time.Millisecond, b.Stop)
-	b.Start() // stops after some delay
+	go b.Start()
+	<-tp.done
+	b.Stop()
+
 	assert.True(t, ok)
 }
 

poller.go

@@ -41,28 +41,20 @@ func NewMiddlewarePoller(original Poller, filter func(*Update) bool) *Middleware
 
 // Poll sieves updates through middleware filter.
 func (p *MiddlewarePoller) Poll(b *Bot, dest chan Update, stop chan struct{}) {
-	cap := 1
-	if p.Capacity > 1 {
-		cap = p.Capacity
+	if p.Capacity < 1 {
+		p.Capacity = 1
 	}
 
-	middle := make(chan Update, cap)
+	middle := make(chan Update, p.Capacity)
 	stopPoller := make(chan struct{})
 
 	go p.Poller.Poll(b, middle, stopPoller)
 
 	for {
 		select {
-
-		// call to stop
 		case <-stop:
-			stopPoller <- struct{}{}
-
-		// poller is done
-		case <-stopPoller:
-			close(stop)
+			close(stopPoller)
 			return
-
 		case upd := <-middle:
 			if p.Filter(&upd) {
 				dest <- upd
@@ -101,7 +93,6 @@ func (p *LongPoller) Poll(b *Bot, dest chan Update, stop chan struct{}) {
 	for {
 		select {
 		case <-stop:
-			close(stop)
 			return
 		default:
 		}

poller_test.go

@@ -1,7 +1,21 @@
 package telebot
 
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
 type testPoller struct {
 	updates chan Update
+	done    chan struct{}
+}
+
+func newTestPoller() *testPoller {
+	return &testPoller{
+		updates: make(chan Update, 1),
+		done:    make(chan struct{}, 1),
+	}
 }
 
 func (p *testPoller) Poll(b *Bot, updates chan Update, stop chan struct{}) {
@@ -10,9 +24,41 @@ func (p *testPoller) Poll(b *Bot, updates chan Update, stop chan struct{}) {
 		case upd := <-p.updates:
 			updates <- upd
 		case <-stop:
-			close(stop)
 			return
 		default:
 		}
 	}
 }
+
+func TestMiddlewarePoller(t *testing.T) {
+	tp := newTestPoller()
+	var ids []int
+
+	b, err := newTestBot()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	b.Poller = NewMiddlewarePoller(tp, func(u *Update) bool {
+		if u.ID > 0 {
+			ids = append(ids, u.ID)
+			return true
+		}
+
+		tp.done <- struct{}{}
+		return false
+	})
+
+	go func() {
+		tp.updates <- Update{ID: 1}
+		tp.updates <- Update{ID: 2}
+		tp.updates <- Update{ID: 0}
+	}()
+
+	go b.Start()
+	<-tp.done
+	b.Stop()
+
+	assert.Contains(t, ids, 1)
+	assert.Contains(t, ids, 2)
+}