Fix for WithTenant() lost orig context

Signed-off-by: Ed Snible <snible@us.ibm.com>
jaegertracing · May 16, 2022 · c076573 · c076573
1 parent 2fd214e
commit c076573
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 1 deletion.
diff --git a/storage/tenant.go b/storage/tenant.go
@@ -26,7 +26,7 @@ const (
 
 // WithTenant creates a Context with a tenant association
 func WithTenant(ctx context.Context, tenant string) context.Context {
-	return context.WithValue(context.Background(), tenantKey, tenant)
+	return context.WithValue(ctx, tenantKey, tenant)
 }
 
 // GetTenant retrieves a tenant associated with a Context

diff --git a/storage/tenant_test.go b/storage/tenant_test.go
@@ -21,11 +21,22 @@ import (
 	"github.com/stretchr/testify/assert"
 )
 
+type testContextKey string
+
 func TestContextTenantHandling(t *testing.T) {
 	ctxWithTenant := WithTenant(context.Background(), "tenant1")
 	assert.Equal(t, "tenant1", GetTenant(ctxWithTenant))
 }
 
+func TestContextPreserved(t *testing.T) {
+	key := testContextKey("expected-key")
+	val := "expected-value"
+	ctxWithValue := context.WithValue(context.Background(), key, val)
+	ctxWithTenant := WithTenant(ctxWithValue, "tenant1")
+	assert.Equal(t, "tenant1", GetTenant(ctxWithTenant))
+	assert.Equal(t, val, ctxWithTenant.Value(key))
+}
+
 func TestNoTenant(t *testing.T) {
 	// If no tenant in context, GetTenant should return the empty string
 	assert.Equal(t, "", GetTenant(context.Background()))