tests: implement poller tests

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)
+}