Update broadcast tests, improve stop mechanism.

Signed-off-by: Valery Piashchynski <[email protected]>

1 files changed, 8 insertions, 2 deletions

diff --git a/plugins/redis/pubsub/pubsub.go b/plugins/redis/pubsub/pubsub.go
index 95a9f6dd..c9ad3d58 100644
--- a/plugins/redis/pubsub/pubsub.go
+++ b/plugins/redis/pubsub/pubsub.go
@@ -172,6 +172,12 @@ func (p *PubSubDriver) Connections(topic string, res map[string]struct{}) {
 }
 
 // Next return next message
-func (p *PubSubDriver) Next() (*pubsub.Message, error) {
-	return p.channel.message(), nil
+func (p *PubSubDriver) Next(ctx context.Context) (*pubsub.Message, error) {
+	const op = errors.Op("redis_driver_next")
+	select {
+	case msg := <-p.channel.message():
+		return msg, nil
+	case <-ctx.Done():
+		return nil, errors.E(op, errors.TimeOut, ctx.Err())
+	}
 }