Function TestDrain

test/drain_test.go:29–75

Drain can be very useful for graceful shutdown of subscribers. Especially queue subscribers.

(t *testing.T)

// Drain can be very useful for graceful shutdown of subscribers.
// Especially queue subscribers.
func TestDrain(t *testing.T) {
	s := RunDefaultServer()
	defer s.Shutdown()
	nc := NewDefaultConnection(t)
	defer nc.Close()

	done := make(chan bool)
	received := int32(0)
	expected := int32(100)

	cb := func(_ *nats.Msg) {
		// Allow this to back up.
		time.Sleep(time.Millisecond)
		rcvd := atomic.AddInt32(&received, 1)
		if rcvd >= expected {
			done <- true
		}
	}

	sub, err := nc.Subscribe("foo", cb)
	if err != nil {
		t.Fatalf("Error creating subscription; %v", err)
	}

	for i := int32(0); i < expected; i++ {
		nc.Publish("foo", []byte("Don't forget about me"))
	}

	// Drain it and make sure we receive all messages.
	sub.Drain()
	if !sub.IsDraining() {
		t.Fatalf("Expected to be draining")
	}
	select {
	case <-done:
		break
	case <-time.After(5 * time.Second):
		r := atomic.LoadInt32(&received)
		if r != expected {
			t.Fatalf("Did not receive all messages: %d of %d", r, expected)
		}
	}
	time.Sleep(100 * time.Millisecond)
	if sub.IsDraining() {
		t.Fatalf("Expected to be done draining")
	}
}

func TestDrainQueueSub(t *testing.T) {
	s := RunDefaultServer()