From 9e8b5fd121e7f45e436bcf59ba854df5df19fd1c Mon Sep 17 00:00:00 2001
From: Krzesimir Nowak <krzesimir@kinvolk.io>
Date: Wed, 21 Aug 2019 17:39:20 +0200
Subject: [PATCH] Implement new interfaces for mock types

---
 mocktracer/mockspan.go        |  5 +++++
 mocktracer/mocktracer.go      | 13 +++++++++++++
 mocktracer/mocktracer_test.go | 23 +++++++++++++++++++++++
 3 files changed, 41 insertions(+)

diff --git a/mocktracer/mockspan.go b/mocktracer/mockspan.go
index 8c7932c..385a149 100644
--- a/mocktracer/mockspan.go
+++ b/mocktracer/mockspan.go
@@ -282,3 +282,8 @@ func (s *MockSpan) SetOperationName(operationName string) opentracing.Span {
 func (s *MockSpan) Tracer() opentracing.Tracer {
 	return s.tracer
 }
+
+// TracerExt1 belongs to the SpanExt1 interface
+func (s *MockSpan) TracerExt1() opentracing.TracerExt1 {
+	return s.tracer
+}
diff --git a/mocktracer/mocktracer.go b/mocktracer/mocktracer.go
index 4533da7..c2c36dd 100644
--- a/mocktracer/mocktracer.go
+++ b/mocktracer/mocktracer.go
@@ -1,6 +1,7 @@
 package mocktracer
 
 import (
+	"context"
 	"sync"
 
 	"github.com/opentracing/opentracing-go"
@@ -59,6 +60,18 @@ func (t *MockTracer) Reset() {
 
 // StartSpan belongs to the Tracer interface.
 func (t *MockTracer) StartSpan(operationName string, opts ...opentracing.StartSpanOption) opentracing.Span {
+	return t.startMockSpan(operationName, opts...)
+}
+
+// StartSpanWithContext belongs to the TracerExt1 interface.
+func (t *MockTracer) StartSpanWithContext(ctx context.Context, operationName string, opts ...opentracing.StartSpanOption) (context.Context, opentracing.SpanExt1) {
+	span := t.startMockSpan(operationName, opts...)
+	ctx = opentracing.ContextWithSpan(ctx, span)
+	return ctx, span
+}
+
+// StartSpan belongs to the Tracer interface.
+func (t *MockTracer) startMockSpan(operationName string, opts ...opentracing.StartSpanOption) *MockSpan {
 	sso := opentracing.StartSpanOptions{}
 	for _, o := range opts {
 		o.Apply(&sso)
diff --git a/mocktracer/mocktracer_test.go b/mocktracer/mocktracer_test.go
index 14c04d8..651b126 100644
--- a/mocktracer/mocktracer_test.go
+++ b/mocktracer/mocktracer_test.go
@@ -1,6 +1,7 @@
 package mocktracer
 
 import (
+	"context"
 	"net/http"
 	"reflect"
 	"sync"
@@ -34,6 +35,28 @@ func TestMockTracer_StartSpan(t *testing.T) {
 	assert.Equal(t, child.ParentID, parent.Context().(MockSpanContext).SpanID)
 }
 
+func TestMockTracer_StartSpanWithContext(t *testing.T) {
+	tracer := New()
+	ctx1, span1 := tracer.StartSpanWithContext(
+		context.Background(),
+		"a",
+		opentracing.Tags(map[string]interface{}{"x": "y"}))
+
+	ctx2, span2 := span1.TracerExt1().StartSpanWithContext(
+		ctx1, "", opentracing.ChildOf(span1.Context()))
+	span2.Finish()
+	span1.Finish()
+	spans := tracer.FinishedSpans()
+	assert.Equal(t, 2, len(spans))
+
+	parent := spans[1]
+	child := spans[0]
+	assert.Equal(t, map[string]interface{}{"x": "y"}, parent.Tags())
+	assert.Equal(t, child.ParentID, parent.Context().(MockSpanContext).SpanID)
+	assert.Equal(t, parent.Context().(MockSpanContext).SpanID, opentracing.SpanFromContext(ctx1).(*MockSpan).Context().(MockSpanContext).SpanID)
+	assert.Equal(t, child.Context().(MockSpanContext).SpanID, opentracing.SpanFromContext(ctx2).(*MockSpan).Context().(MockSpanContext).SpanID)
+}
+
 func TestMockSpan_SetOperationName(t *testing.T) {
 	tracer := New()
 	span := tracer.StartSpan("")