From 818c2d6f359bba9ea8e5bbf1b447bcde92a4c176 Mon Sep 17 00:00:00 2001
From: Ken Hibino <ken.hibino7@gmail.com>
Date: Mon, 7 Sep 2020 10:54:42 -0700
Subject: [PATCH] Add GetQueueName helper to extract queue name from context

---
 context.go      | 13 +++++++++++++
 context_test.go | 16 ++++++++++++++--
 2 files changed, 27 insertions(+), 2 deletions(-)

diff --git a/context.go b/context.go
index 1ed6634..4710bbd 100644
--- a/context.go
+++ b/context.go
@@ -16,6 +16,7 @@ type taskMetadata struct {
 	id         string
 	maxRetry   int
 	retryCount int
+	qname      string
 }
 
 // ctxKey type is unexported to prevent collisions with context keys defined in
@@ -32,6 +33,7 @@ func createContext(msg *base.TaskMessage, deadline time.Time) (context.Context,
 		id:         msg.ID.String(),
 		maxRetry:   msg.Retry,
 		retryCount: msg.Retried,
+		qname:      msg.Queue,
 	}
 	ctx := context.WithValue(context.Background(), metadataCtxKey, metadata)
 	return context.WithDeadline(ctx, deadline)
@@ -72,3 +74,14 @@ func GetMaxRetry(ctx context.Context) (n int, ok bool) {
 	}
 	return metadata.maxRetry, true
 }
+
+// GetQueueName extracts queue name from a context, if any.
+//
+// Return value qname indicates which queue the task was pulled from.
+func GetQueueName(ctx context.Context) (qname string, ok bool) {
+	metadata, ok := ctx.Value(metadataCtxKey).(taskMetadata)
+	if !ok {
+		return "", false
+	}
+	return metadata.qname, true
+}
diff --git a/context_test.go b/context_test.go
index af11d1d..305cd19 100644
--- a/context_test.go
+++ b/context_test.go
@@ -92,8 +92,9 @@ func TestGetTaskMetadataFromContext(t *testing.T) {
 		desc string
 		msg  *base.TaskMessage
 	}{
-		{"with zero retried message", &base.TaskMessage{Type: "something", ID: uuid.New(), Retry: 25, Retried: 0, Timeout: 1800}},
-		{"with non-zero retried message", &base.TaskMessage{Type: "something", ID: uuid.New(), Retry: 10, Retried: 5, Timeout: 1800}},
+		{"with zero retried message", &base.TaskMessage{Type: "something", ID: uuid.New(), Retry: 25, Retried: 0, Timeout: 1800, Queue: "default"}},
+		{"with non-zero retried message", &base.TaskMessage{Type: "something", ID: uuid.New(), Retry: 10, Retried: 5, Timeout: 1800, Queue: "default"}},
+		{"with custom queue name", &base.TaskMessage{Type: "something", ID: uuid.New(), Retry: 25, Retried: 0, Timeout: 1800, Queue: "custom"}},
 	}
 
 	for _, tc := range tests {
@@ -123,6 +124,14 @@ func TestGetTaskMetadataFromContext(t *testing.T) {
 		if ok && maxRetry != tc.msg.Retry {
 			t.Errorf("%s: GetMaxRetry(ctx) returned n == %d want %d", tc.desc, maxRetry, tc.msg.Retry)
 		}
+
+		qname, ok := GetQueueName(ctx)
+		if !ok {
+			t.Errorf("%s: GetQueueName(ctx) returned ok == false", tc.desc)
+		}
+		if ok && qname != tc.msg.Queue {
+			t.Errorf("%s: GetQueueName(ctx) returned qname == %q, want %q", tc.desc, qname, tc.msg.Queue)
+		}
 	}
 }
 
@@ -144,5 +153,8 @@ func TestGetTaskMetadataFromContextError(t *testing.T) {
 		if _, ok := GetMaxRetry(tc.ctx); ok {
 			t.Errorf("%s: GetMaxRetry(ctx) returned ok == true", tc.desc)
 		}
+		if _, ok := GetQueueName(tc.ctx); ok {
+			t.Errorf("%s: GetQueueName(ctx) returned ok == true", tc.desc)
+		}
 	}
 }