Tidy up

pubsub/subscriber.go

···
 			return
 		}
 
-		err := s.readMessage(consumer.msgs)
+		err := s.readMessage(ctx, consumer.msgs)
 		if err != nil {
 			consumer.Err = err
 			return
···
 	}
 }
 
-func (s *Subscriber) readMessage(msgChan chan *Message) error {
-	// var msg *Message
+func (s *Subscriber) readMessage(ctx context.Context, msgChan chan *Message) error {
 	op := func(conn net.Conn) error {
 		err := s.conn.SetReadDeadline(time.Now().Add(time.Second))
 		if err != nil {
···
 
 		msgChan <- msg
 
-		ack := <-msg.ack
+		var ack bool
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		case ack = <-msg.ack:
+		}
+		//ack := <-msg.ack
 
 		ackMessage := server.Nack
 		if ack {

pubsub/subscriber_test.go

···
 	err = publisher.PublishMessage(msg)
 	require.NoError(t, err)
 
+	time.Sleep(time.Second)
 	cancel()
 
 	select {
···
 	}
 
 	// give the consumer some time to read the messages -- TODO: make better!
-	time.Sleep(time.Millisecond * 500)
+	time.Sleep(time.Second)
 	cancel()
 
 	select {
···
 	require.NoError(t, err)
 
 	// give the consumer some time to read the messages -- TODO: make better!
-	time.Sleep(time.Millisecond * 500)
+	time.Sleep(time.Second)
 	cancel()
 
 	select {