Address comments

uber · Aug 3, 2017 · 53f76ff · 53f76ff
1 parent 86fcba8
commit 53f76ff
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 2 deletions.
diff --git a/context_header.go b/context_header.go
@@ -115,7 +115,7 @@ func WrapWithHeaders(ctx context.Context, headers map[string]string) ContextWith
 	return headerCtx{Context: newCtx}
 }
 
-// WithoutHeaders removes any TChannel headers from the given context.
+// WithoutHeaders hides any TChannel headers from the given context.
 func WithoutHeaders(ctx context.Context) context.Context {
 	return withoutHeaders{ctx}
 }

diff --git a/context_internal_test.go b/context_internal_test.go
@@ -56,11 +56,19 @@ func TestCurrentSpan(t *testing.T) {
 	assert.NotEqual(t, uint64(0), span.TraceID(), "mock tracer is now Zipkin-compatible")
 }
 
-func TestContextWithoutHeaders(t *testing.T) {
+func TestContextWithoutHeadersKeyHeaders(t *testing.T) {
 	ctx := WrapWithHeaders(context.Background(), map[string]string{"k1": "v1"})
 	assert.Equal(t, map[string]string{"k1": "v1"}, ctx.Headers())
 	ctx2 := WithoutHeaders(ctx)
 	assert.Nil(t, ctx2.Value(contextKeyHeaders))
 	_, ok := ctx2.(ContextWithHeaders)
 	assert.False(t, ok)
 }
+
+func TestContextWithoutHeadersKeyTChannel(t *testing.T) {
+	ctx, _ := NewContextBuilder(time.Second).SetShardKey("s1").Build()
+	ctx2 := WithoutHeaders(ctx)
+	assert.Nil(t, ctx2.Value(contextKeyTChannel))
+	_, ok := ctx2.(ContextWithHeaders)
+	assert.False(t, ok)
+}