chore(queue): test Capacity and Usage

Signed-off-by: Bo-Yi Wu <[email protected]>

Author: appleboy

queue_test.go

@@ -2,10 +2,19 @@ package queue
 
 import (
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 )
 
+type mockMessage struct {
+	message string
+}
+
+func (m mockMessage) Bytes() []byte {
+	return []byte(m.message)
+}
+
 func TestNewQueue(t *testing.T) {
 	q, err := NewQueue()
 	assert.Error(t, err)
@@ -18,3 +27,69 @@ func TestNewQueue(t *testing.T) {
 	assert.NoError(t, err)
 	assert.NotNil(t, q)
 }
+
+func TestWorkerNum(t *testing.T) {
+	w := &queueWorker{
+		messages: make(chan QueuedMessage, 100),
+	}
+	q, err := NewQueue(
+		WithWorker(w),
+		WithWorkerCount(2),
+	)
+	assert.NoError(t, err)
+	assert.NotNil(t, q)
+
+	q.Start()
+	q.Start()
+	time.Sleep(20 * time.Millisecond)
+	assert.Equal(t, 4, q.Workers())
+	q.Shutdown()
+	q.Wait()
+}
+
+func TestShtdonwOnce(t *testing.T) {
+	w := &queueWorker{
+		messages: make(chan QueuedMessage, 100),
+	}
+	q, err := NewQueue(
+		WithWorker(w),
+		WithWorkerCount(2),
+	)
+	assert.NoError(t, err)
+	assert.NotNil(t, q)
+
+	q.Start()
+	time.Sleep(20 * time.Millisecond)
+	assert.Equal(t, 2, q.Workers())
+	q.Shutdown()
+	// don't panic here
+	q.Shutdown()
+	q.Wait()
+	assert.Equal(t, 0, q.Workers())
+}
+
+func TestWorkerStatus(t *testing.T) {
+	m := mockMessage{
+		message: "foobar",
+	}
+	w := &queueWorker{
+		messages: make(chan QueuedMessage, 100),
+	}
+	q, err := NewQueue(
+		WithWorker(w),
+		WithWorkerCount(2),
+	)
+	assert.NoError(t, err)
+	assert.NotNil(t, q)
+
+	q.Queue(m)
+	q.Queue(m)
+	q.Queue(m)
+	q.Queue(m)
+	assert.Equal(t, 100, q.Capacity())
+	assert.Equal(t, 4, q.Usage())
+	q.Start()
+	time.Sleep(20 * time.Millisecond)
+	q.Shutdown()
+	q.Wait()
+}

worker.go

@@ -1,5 +1,11 @@
 package queue
 
+import (
+	"errors"
+	"log"
+	"time"
+)
+
 // Worker interface
 type Worker interface {
 	BeforeRun() error
@@ -28,3 +34,33 @@ func (w *emptyWorker) Shutdown() error               { return nil }
 func (w *emptyWorker) Queue(job QueuedMessage) error { return nil }
 func (w *emptyWorker) Capacity() int                 { return 0 }
 func (w *emptyWorker) Usage() int                    { return 0 }
+
+type queueWorker struct {
+	messages chan QueuedMessage
+}
+
+func (w *queueWorker) BeforeRun() error { return nil }
+func (w *queueWorker) AfterRun() error  { return nil }
+func (w *queueWorker) Run(chan struct{}) error {
+	for msg := range w.messages {
+		log.Println("got message", msg)
+		time.Sleep(100 * time.Millisecond)
+	}
+	return nil
+}
+
+func (w *queueWorker) Shutdown() error {
+	close(w.messages)
+	return nil
+}
+
+func (w *queueWorker) Queue(job QueuedMessage) error {
+	select {
+	case w.messages <- job:
+		return nil
+	default:
+		return errors.New("max capacity reached")
+	}
+}
+func (w *queueWorker) Capacity() int { return cap(w.messages) }
+func (w *queueWorker) Usage() int    { return len(w.messages) }