Ensure events get handled in the right order (#47)

Resolves #40

consumer/consumer.go

@@ -3,6 +3,7 @@ package consumer
 import (
 	"context"
 	"fmt"
+	"sync"
 
 	"github.com/bketelsen/logr"
 	"github.com/corvus-ch/rabbitmq-cli-consumer/delivery"
@@ -18,6 +19,9 @@ type Consumer struct {
 	Processor  processor.Processor
 	Log        logr.Logger
 	canceled   bool
+
+	// wg is used to ensure NotifyClose() gets handled before the consumer exits.
+	wg sync.WaitGroup
 }
 
 // New creates a new consumer instance. The setup of the amqp connection and channel is expected to be done by the
@@ -102,6 +106,7 @@ func (c *Consumer) consume(msgs <-chan amqp.Delivery, done chan error) {
 			return
 		}
 	}
+	c.wg.Wait()
 	done <- nil
 }
 
@@ -129,13 +134,15 @@ func (c *Consumer) Close() error {
 // The chan provided will be closed when the Channel is closed and on a Graceful close, no error will be sent.
 func (c *Consumer) NotifyClose(receiver chan error) chan error {
 	if c.Channel != nil {
+		c.wg.Add(1)
 		realChan := make(chan *amqp.Error)
 		c.Channel.NotifyClose(realChan)
 
 		go func() {
 			for {
 				err, ok := <-realChan
 				if !ok {
+					c.wg.Done()
 					return
 				}
 				receiver <- err

integration_test.go

@@ -167,6 +167,36 @@ func TestEndToEnd(t *testing.T) {
 	})
 }
 
+var closeTests = []struct {
+	name      string
+	closeArgs []string
+}{
+	{"connection", []string{"exec", "-T", "rabbitmq", "/bin/bash", "-c", "until rabbitmqadmin list connections | grep -v 'No items' > /dev/null; do sleep 1; done; eval $(rabbitmqadmin list connections -f kvp); rabbitmqadmin close connection name=\"${name}\""}},
+	{"shutdown", []string{"stop"}},
+}
+
+func TestConnectionClose(t *testing.T) {
+	for _, test := range closeTests {
+		t.Run(test.name, func(t *testing.T) {
+			conn, ch := prepare(t)
+
+			args := []string{"-V", "-no-datetime", "-e", command, "-c", "fixtures/default.conf"}
+			cmd, stdout, _ := startConsumer(t, []string{}, args...)
+			waitForOutput(t, stdout, "Waiting for messages...")
+
+			conn.Close()
+			ch.Close()
+
+			stop := exec.Command("docker-compose", test.closeArgs...)
+			if err := stop.Run(); err != nil {
+				t.Fatalf("failed to close connection/shutdown server: %v", err)
+			}
+
+			assert.EqualError(t, cmd.Wait(), "exit status 10")
+		})
+	}
+}
+
 func assertOutput(t *testing.T, stdout, stderr *bytes.Buffer) {
 	goldie.Assert(t, t.Name()+"Output", bytes.Trim(stdout.Bytes(), "\x00"))
 	goldie.Assert(t, t.Name()+"Error", bytes.Trim(stderr.Bytes(), "\x00"))