func TestTraceServer(t *testing.T) {
	tracer := mocktracer.New()

	// Initialize the ctx with a nameless Span.
	contextSpan := tracer.StartSpan("").(*mocktracer.MockSpan)
	ctx := opentracing.ContextWithSpan(context.Background(), contextSpan)

	var innerEndpoint endpoint.Endpoint
	innerEndpoint = func(context.Context, interface{}) (interface{}, error) {
		return struct{}{}, nil
	}
	tracedEndpoint := kitot.TraceServer(tracer, "testOp")(innerEndpoint)
	if _, err := tracedEndpoint(ctx, struct{}{}); err != nil {
		t.Fatal(err)
	}
	if want, have := 1, len(tracer.FinishedSpans); want != have {
		t.Fatalf("Want %v span(s), found %v", want, have)
	}

	endpointSpan := tracer.FinishedSpans[0]
	// Test that the op name is updated
	if want, have := "testOp", endpointSpan.OperationName; want != have {
		t.Fatalf("Want %q, have %q", want, have)
	}
	// ... and that the ID is unmodified.
	if want, have := contextSpan.SpanID, endpointSpan.SpanID; want != have {
		t.Errorf("Want SpanID %q, have %q", want, have)
	}
}

func TestRPCServerOption(t *testing.T) {
	tracer := mocktracer.New()
	parent := tracer.StartSpan("my-trace")
	parent.SetBaggageItem("bag", "gage")

	carrier := opentracing.HTTPHeadersCarrier{}
	err := tracer.Inject(parent.Context(), opentracing.HTTPHeaders, carrier)
	if err != nil {
		t.Fatal(err)
	}

	parCtx, err := tracer.Extract(opentracing.HTTPHeaders, carrier)
	if err != nil {
		t.Fatal(err)
	}

	tracer.StartSpan("my-child", ext.RPCServerOption(parCtx)).Finish()

	rawSpan := tracer.FinishedSpans()[0]
	assert.Equal(t, map[string]interface{}{
		"span.kind": ext.SpanKindRPCServerEnum,
	}, rawSpan.Tags())
	assert.Equal(t, map[string]string{
		"bag": "gage",
	}, rawSpan.Context().(mocktracer.MockSpanContext).Baggage)
}

func TestTraceServer(t *testing.T) {
	tracer := mocktracer.New()

	// Initialize the ctx with a nameless Span.
	contextSpan := tracer.StartSpan("").(*mocktracer.MockSpan)
	ctx := opentracing.ContextWithSpan(context.Background(), contextSpan)

	tracedEndpoint := kitot.TraceServer(tracer, "testOp")(endpoint.Nop)
	if _, err := tracedEndpoint(ctx, struct{}{}); err != nil {
		t.Fatal(err)
	}

	finishedSpans := tracer.FinishedSpans()
	if want, have := 1, len(finishedSpans); want != have {
		t.Fatalf("Want %v span(s), found %v", want, have)
	}

	// Test that the op name is updated
	endpointSpan := finishedSpans[0]
	if want, have := "testOp", endpointSpan.OperationName; want != have {
		t.Fatalf("Want %q, have %q", want, have)
	}
	contextContext := contextSpan.Context().(mocktracer.MockSpanContext)
	endpointContext := endpointSpan.Context().(mocktracer.MockSpanContext)
	// ...and that the ID is unmodified.
	if want, have := contextContext.SpanID, endpointContext.SpanID; want != have {
		t.Errorf("Want SpanID %q, have %q", want, have)
	}
}

func TestPeerTags(t *testing.T) {
	if ext.PeerService != "peer.service" {
		t.Fatalf("Invalid PeerService %v", ext.PeerService)
	}
	tracer := mocktracer.New()
	span := tracer.StartSpan("my-trace")
	ext.PeerService.Set(span, "my-service")
	ext.PeerHostname.Set(span, "my-hostname")
	ext.PeerHostIPv4.Set(span, uint32(127<<24|1))
	ext.PeerHostIPv6.Set(span, "::")
	ext.PeerPort.Set(span, uint16(8080))
	ext.SamplingPriority.Set(span, uint16(1))
	ext.SpanKind.Set(span, ext.SpanKindRPCServerEnum)
	ext.SpanKindRPCClient.Set(span)
	span.Finish()

	rawSpan := tracer.FinishedSpans()[0]
	assert.Equal(t, map[string]interface{}{
		"peer.service":  "my-service",
		"peer.hostname": "my-hostname",
		"peer.ipv4":     uint32(127<<24 | 1),
		"peer.ipv6":     "::",
		"peer.port":     uint16(8080),
		"span.kind":     ext.SpanKindRPCClientEnum,
	}, rawSpan.Tags())
	assert.True(t, span.Context().(mocktracer.MockSpanContext).Sampled)
	ext.SamplingPriority.Set(span, uint16(0))
	assert.False(t, span.Context().(mocktracer.MockSpanContext).Sampled)
}

func TestTraceClient(t *testing.T) {
	tracer := mocktracer.New()

	// Initialize the ctx with a parent Span.
	parentSpan := tracer.StartSpan("parent").(*mocktracer.MockSpan)
	defer parentSpan.Finish()
	ctx := opentracing.ContextWithSpan(context.Background(), parentSpan)

	tracedEndpoint := kitot.TraceClient(tracer, "testOp")(endpoint.Nop)
	if _, err := tracedEndpoint(ctx, struct{}{}); err != nil {
		t.Fatal(err)
	}

	// tracedEndpoint created a new Span.
	finishedSpans := tracer.FinishedSpans()
	if want, have := 1, len(finishedSpans); want != have {
		t.Fatalf("Want %v span(s), found %v", want, have)
	}

	endpointSpan := finishedSpans[0]
	if want, have := "testOp", endpointSpan.OperationName; want != have {
		t.Fatalf("Want %q, have %q", want, have)
	}

	parentContext := parentSpan.Context().(mocktracer.MockSpanContext)
	endpointContext := parentSpan.Context().(mocktracer.MockSpanContext)

	// ... and that the parent ID is set appropriately.
	if want, have := parentContext.SpanID, endpointContext.SpanID; want != have {
		t.Errorf("Want ParentID %q, have %q", want, have)
	}
}

func TestExtractInboundSpanWithZipkinTracer(t *testing.T) {
	tracer := mocktracer.New()
	callReq := new(callReq)
	callReq.Tracing = Span{traceID: 1, spanID: 2, flags: 1}
	callReq.Headers = transportHeaders{
		ArgScheme:  string(JSON),
		CallerName: "caller",
	}
	c := Connection{
		channelConnectionCommon: channelConnectionCommon{tracer: tracer},
		remotePeerInfo:          PeerInfo{HostPort: "host:123"},
	}
	c.parseRemotePeerAddress()

	// fail to extract with zipkin format, as MockTracer does not support it
	assert.Nil(t, c.extractInboundSpan(callReq), "zipkin format not available")

	// add zipkin format extractor and try again
	tracer.RegisterExtractor(zipkinSpanFormat, new(zipkinExtractor))
	span := c.extractInboundSpan(callReq)
	require.NotNil(t, span, "zipkin format available")

	// validate the extracted span was correctly populated
	s1, ok := span.(*mocktracer.MockSpan)
	require.True(t, ok)
	assert.Equal(t, 1, s1.SpanContext.TraceID)
	assert.Equal(t, 2, s1.ParentID)
	assert.True(t, s1.SpanContext.Sampled)
	assert.Equal(t, "", s1.OperationName, "operation name unknown initially")
	assert.Equal(t, string(JSON), s1.Tag("as"))
	assert.Equal(t, "caller", s1.Tag(string(ext.PeerService)))
	assert.Equal(t, "host", s1.Tag(string(ext.PeerHostname)))
	assert.Equal(t, uint16(123), s1.Tag(string(ext.PeerPort)))

	// start a temporary span so that we can populate headers with baggage
	tempSpan := tracer.StartSpan("test")
	tempSpan.SetBaggageItem("x", "y")
	headers := make(map[string]string)
	carrier := tracingHeadersCarrier(headers)
	err := tracer.Inject(tempSpan.Context(), opentracing.TextMap, carrier)
	assert.NoError(t, err)

	// run the public ExtractInboundSpan method with application headers
	inCall := &InboundCall{
		response: &InboundCallResponse{
			span: span,
		},
	}
	ctx := context.Background()
	ctx2 := ExtractInboundSpan(ctx, inCall, headers, tracer)
	span = opentracing.SpanFromContext(ctx2)
	s2, ok := span.(*mocktracer.MockSpan)
	require.True(t, ok)
	assert.Equal(t, s1, s2, "should be the same span started previously")
	assert.Equal(t, "y", s2.BaggageItem("x"), "baggage should've been added")
}

func TestServer(t *testing.T) {
	tracer := mocktracer.New()

	s := &Server{HostPort: "127.0.0.1:0", Tracer: tracer}
	err := s.Start()
	require.NoError(t, err)
	defer s.Close()

	assert.Equal(t, tracer, s.Ch.Tracer())
	port, err := strconv.Atoi(s.Port())
	assert.NoError(t, err)
	assert.True(t, port > 0)
}

func (s *PropagationTestSuite) runWithMockTracer(t *testing.T) {
	mockTracer := mocktracer.New()
	s.runWithTracer(t, tracerChoice{
		tracerType: Mock,
		tracer:     mockTracer,
		spansRecorded: func() int {
			return len(MockTracerSampledSpans(mockTracer))
		},
		resetSpans: func() {
			mockTracer.Reset()
		},
	})
}

func TestTraceHTTPRequestRoundtrip(t *testing.T) {
	logger := log.NewNopLogger()
	tracer := mocktracer.New()

	// Initialize the ctx with a Span to inject.
	beforeSpan := tracer.StartSpan("to_inject").(*mocktracer.MockSpan)
	defer beforeSpan.Finish()
	beforeSpan.SetBaggageItem("baggage", "check")
	beforeCtx := opentracing.ContextWithSpan(context.Background(), beforeSpan)

	toHTTPFunc := kitot.ToHTTPRequest(tracer, logger)
	req, _ := http.NewRequest("GET", "http://test.biz/url", nil)
	// Call the RequestFunc.
	afterCtx := toHTTPFunc(beforeCtx, req)

	// The Span should not have changed.
	afterSpan := opentracing.SpanFromContext(afterCtx)
	if beforeSpan != afterSpan {
		t.Errorf("Should not swap in a new span")
	}

	// No spans should have finished yet.
	finishedSpans := tracer.FinishedSpans()
	if want, have := 0, len(finishedSpans); want != have {
		t.Errorf("Want %v span(s), found %v", want, have)
	}

	// Use FromHTTPRequest to verify that we can join with the trace given a req.
	fromHTTPFunc := kitot.FromHTTPRequest(tracer, "joined", logger)
	joinCtx := fromHTTPFunc(afterCtx, req)
	joinedSpan := opentracing.SpanFromContext(joinCtx).(*mocktracer.MockSpan)

	joinedContext := joinedSpan.Context().(mocktracer.MockSpanContext)
	beforeContext := beforeSpan.Context().(mocktracer.MockSpanContext)

	if joinedContext.SpanID == beforeContext.SpanID {
		t.Error("SpanID should have changed", joinedContext.SpanID, beforeContext.SpanID)
	}

	// Check that the parent/child relationship is as expected for the joined span.
	if want, have := beforeContext.SpanID, joinedSpan.ParentID; want != have {
		t.Errorf("Want ParentID %q, have %q", want, have)
	}
	if want, have := "joined", joinedSpan.OperationName; want != have {
		t.Errorf("Want %q, have %q", want, have)
	}
	if want, have := "check", joinedSpan.BaggageItem("baggage"); want != have {
		t.Errorf("Want %q, have %q", want, have)
	}
}

func TestMiscTags(t *testing.T) {
	tracer := mocktracer.New()
	span := tracer.StartSpan("my-trace")
	ext.Component.Set(span, "my-awesome-library")
	ext.SamplingPriority.Set(span, 1)
	ext.Error.Set(span, true)

	span.Finish()

	rawSpan := tracer.FinishedSpans()[0]
	assert.Equal(t, map[string]interface{}{
		"component": "my-awesome-library",
		"error":     true,
	}, rawSpan.Tags())
}

func TestHTTPTags(t *testing.T) {
	tracer := mocktracer.New()
	span := tracer.StartSpan("my-trace", ext.SpanKindRPCServer)
	ext.HTTPUrl.Set(span, "test.biz/uri?protocol=false")
	ext.HTTPMethod.Set(span, "GET")
	ext.HTTPStatusCode.Set(span, 301)
	span.Finish()

	rawSpan := tracer.FinishedSpans()[0]
	assert.Equal(t, map[string]interface{}{
		"http.url":         "test.biz/uri?protocol=false",
		"http.method":      "GET",
		"http.status_code": uint16(301),
		"span.kind":        ext.SpanKindRPCServerEnum,
	}, rawSpan.Tags())
}

func TestCurrentSpan(t *testing.T) {
	ctx := context.Background()
	span := CurrentSpan(ctx)
	require.NotNil(t, span, "CurrentSpan() should always return something")

	tracer := mocktracer.New()
	sp := tracer.StartSpan("test")
	ctx = opentracing.ContextWithSpan(ctx, sp)
	span = CurrentSpan(ctx)
	require.NotNil(t, span, "CurrentSpan() should always return something")
	assert.EqualValues(t, 0, span.TraceID(), "mock tracer is not Zipkin-compatible")

	tracer.RegisterInjector(zipkinSpanFormat, new(zipkinInjector))
	span = CurrentSpan(ctx)
	require.NotNil(t, span, "CurrentSpan() should always return something")
	assert.NotEqual(t, uint64(0), span.TraceID(), "mock tracer is now Zipkin-compatible")
}

func TestTraceGRPCRequestRoundtrip(t *testing.T) {
	tracer := mocktracer.New()

	// Initialize the ctx with a Span to inject.
	beforeSpan := tracer.StartSpan("to_inject").(*mocktracer.MockSpan)
	defer beforeSpan.Finish()
	beforeSpan.SetBaggageItem("baggage", "check")
	beforeCtx := opentracing.ContextWithSpan(context.Background(), beforeSpan)

	var toGRPCFunc grpc.RequestFunc = kitot.ToGRPCRequest(tracer, nil)
	md := metadata.Pairs()
	// Call the RequestFunc.
	afterCtx := toGRPCFunc(beforeCtx, &md)

	// The Span should not have changed.
	afterSpan := opentracing.SpanFromContext(afterCtx)
	if beforeSpan != afterSpan {
		t.Errorf("Should not swap in a new span")
	}

	// No spans should have finished yet.
	if want, have := 0, len(tracer.FinishedSpans); want != have {
		t.Errorf("Want %v span(s), found %v", want, have)
	}

	// Use FromGRPCRequest to verify that we can join with the trace given MD.
	var fromGRPCFunc grpc.RequestFunc = kitot.FromGRPCRequest(tracer, "joined", nil)
	joinCtx := fromGRPCFunc(afterCtx, &md)
	joinedSpan := opentracing.SpanFromContext(joinCtx).(*mocktracer.MockSpan)

	if joinedSpan.SpanID == beforeSpan.SpanID {
		t.Error("SpanID should have changed", joinedSpan.SpanID, beforeSpan.SpanID)
	}

	// Check that the parent/child relationship is as expected for the joined span.
	if want, have := beforeSpan.SpanID, joinedSpan.ParentID; want != have {
		t.Errorf("Want ParentID %q, have %q", want, have)
	}
	if want, have := "joined", joinedSpan.OperationName; want != have {
		t.Errorf("Want %q, have %q", want, have)
	}
	if want, have := "check", joinedSpan.BaggageItem("baggage"); want != have {
		t.Errorf("Want %q, have %q", want, have)
	}
}

func TestTracingInjectorExtractor(t *testing.T) {
	tracer := mocktracer.New()
	tracer.RegisterInjector(zipkinSpanFormat, new(zipkinInjector))
	tracer.RegisterExtractor(zipkinSpanFormat, new(zipkinExtractor))

	sp := tracer.StartSpan("x")
	var injectable injectableSpan
	err := tracer.Inject(sp.Context(), zipkinSpanFormat, &injectable)
	require.NoError(t, err)

	tsp := Span(injectable)
	assert.NotEqual(t, uint64(0), tsp.TraceID())
	assert.NotEqual(t, uint64(0), tsp.SpanID())

	sp2, err := tracer.Extract(zipkinSpanFormat, &tsp)
	require.NoError(t, err)
	require.NotNil(t, sp2)
}

func TestTChannelTracerDepth2(t *testing.T) {
	tracer := mocktracer.New()
	dispatcher := createTChannelDispatcher(tracer, t)

	client := json.New(dispatcher.Channel("yarpc-test"))
	handler := handler{client: client, t: t}
	handler.register(dispatcher)

	require.NoError(t, dispatcher.Start())
	defer dispatcher.Stop()

	ctx, cancel := handler.createContextWithBaggage(tracer)
	defer cancel()

	err := handler.echoEcho(ctx)
	assert.NoError(t, err)

	AssertDepth2Spans(t, tracer)
}

func TestTraceClientNoContextSpan(t *testing.T) {
	tracer := mocktracer.New()

	// Empty/background context.
	tracedEndpoint := kitot.TraceClient(tracer, "testOp")(endpoint.Nop)
	if _, err := tracedEndpoint(context.Background(), struct{}{}); err != nil {
		t.Fatal(err)
	}

	// tracedEndpoint created a new Span.
	finishedSpans := tracer.FinishedSpans()
	if want, have := 1, len(finishedSpans); want != have {
		t.Fatalf("Want %v span(s), found %v", want, have)
	}

	endpointSpan := finishedSpans[0]
	if want, have := "testOp", endpointSpan.OperationName; want != have {
		t.Fatalf("Want %q, have %q", want, have)
	}
}

func TestChannelTracerMethod(t *testing.T) {
	mockTracer := mocktracer.New()
	ch, err := NewChannel("svc", &ChannelOptions{
		Tracer: mockTracer,
	})
	require.NoError(t, err)
	defer ch.Close()
	assert.Equal(t, mockTracer, ch.Tracer(), "expecting tracer passed at initialization")

	ch, err = NewChannel("svc", &ChannelOptions{})
	require.NoError(t, err)
	defer ch.Close()
	assert.EqualValues(t, opentracing.GlobalTracer(), ch.Tracer(), "expecting default tracer")

	// because ch.Tracer() function is doing dynamic lookup, we can change global tracer
	origTracer := opentracing.GlobalTracer()
	defer opentracing.InitGlobalTracer(origTracer)

	opentracing.InitGlobalTracer(mockTracer)
	assert.Equal(t, mockTracer, ch.Tracer(), "expecting tracer set as global tracer")
}

func TestSetPeerHostPort(t *testing.T) {
	tracer := mocktracer.New()

	ipv6 := []byte{1, 2, 3, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 15, 16}
	assert.Equal(t, net.IPv6len, len(ipv6))
	ipv6hostPort := fmt.Sprintf("[%v]:789", net.IP(ipv6))

	tests := []struct {
		hostPort    string
		wantHostTag string
		wantHost    interface{}
		wantPort    uint16
	}{
		{"adhoc123:bad-port", "peer.hostname", "adhoc123", 0},
		{"adhoc123", "peer.hostname", "adhoc123", 0},
		{"ip123.uswest.aws.com:765", "peer.hostname", "ip123.uswest.aws.com", 765},
		{"localhost:123", "peer.ipv4", uint32(127<<24 | 1), 123},
		{"10.20.30.40:321", "peer.ipv4", uint32(10<<24 | 20<<16 | 30<<8 | 40), 321},
		{ipv6hostPort, "peer.ipv6", "102:300::f10", 789},
	}

	for i, test := range tests {
		span := tracer.StartSpan("x")
		c := &Connection{
			remotePeerInfo: PeerInfo{
				HostPort: test.hostPort,
			},
		}
		c.parseRemotePeerAddress()
		c.setPeerHostPort(span)
		span.Finish()
		rawSpan := tracer.FinishedSpans()[i]
		assert.Equal(t, test.wantHost, rawSpan.Tag(test.wantHostTag), "test %+v", test)
		if test.wantPort != 0 {
			assert.Equal(t, test.wantPort, rawSpan.Tag(string(ext.PeerPort)), "test %+v", test)
		} else {
			assert.Nil(t, rawSpan.Tag(string(ext.PeerPort)), "test %+v", test)
		}
	}
}

func TestTraceServerNoContextSpan(t *testing.T) {
	tracer := mocktracer.New()

	var innerEndpoint endpoint.Endpoint
	innerEndpoint = func(context.Context, interface{}) (interface{}, error) {
		return struct{}{}, nil
	}
	tracedEndpoint := kitot.TraceServer(tracer, "testOp")(innerEndpoint)
	// Empty/background context:
	if _, err := tracedEndpoint(context.Background(), struct{}{}); err != nil {
		t.Fatal(err)
	}
	// tracedEndpoint created a new Span:
	if want, have := 1, len(tracer.FinishedSpans); want != have {
		t.Fatalf("Want %v span(s), found %v", want, have)
	}

	endpointSpan := tracer.FinishedSpans[0]
	if want, have := "testOp", endpointSpan.OperationName; want != have {
		t.Fatalf("Want %q, have %q", want, have)
	}
}

// Per https://github.com/uber/tchannel-go/issues/505, concurrent client calls
// made with the same shared map used as headers were causing panic due to
// concurrent writes to the map when injecting tracing headers.
func TestReusableHeaders(t *testing.T) {
	opts := &testutils.ChannelOpts{
		ChannelOptions: ChannelOptions{Tracer: mocktracer.New()},
	}
	WithVerifiedServer(t, opts, func(ch *Channel, hostPort string) {
		jsonHandler := &JSONHandler{TraceHandler: testtracing.TraceHandler{Ch: ch}, t: t}
		json.Register(ch, json.Handlers{"call": jsonHandler.callJSON}, jsonHandler.onError)

		span := ch.Tracer().StartSpan("client")
		traceID := span.(*mocktracer.MockSpan).SpanContext.TraceID // for validation
		ctx := opentracing.ContextWithSpan(context.Background(), span)

		sharedHeaders := map[string]string{"life": "42"}
		ctx, cancel := NewContextBuilder(2 * time.Second).
			SetHeaders(sharedHeaders).
			SetParentContext(ctx).
			Build()
		defer cancel()

		peer := ch.Peers().GetOrAdd(ch.PeerInfo().HostPort)

		var wg sync.WaitGroup
		for i := 0; i < 42; i++ {
			wg.Add(1)
			go func() {
				defer wg.Done()
				var response testtracing.TracingResponse
				err := json.CallPeer(json.Wrap(ctx), peer, ch.ServiceName(),
					"call", &testtracing.TracingRequest{}, &response)
				assert.NoError(t, err, "json.Call failed")
				assert.EqualValues(t, traceID, response.TraceID, "traceID must match")
			}()
		}
		wg.Wait()
		assert.Equal(t, map[string]string{"life": "42"}, sharedHeaders, "headers unchanged")
	})
}