From 79b602f6ab1bfd3ea752abb61b4a4a45af63b2ee Mon Sep 17 00:00:00 2001
From: Kofo Okesola <okesolakofo@gmail.com>
Date: Mon, 31 Jul 2023 09:34:13 +0100
Subject: [PATCH] added context with span

---
 trace/span.go      |  5 +++++
 trace/span_test.go | 14 ++++++++++++++
 2 files changed, 19 insertions(+)

diff --git a/trace/span.go b/trace/span.go
index 95edf64..1cfe6fd 100644
--- a/trace/span.go
+++ b/trace/span.go
@@ -18,6 +18,11 @@ func SpanFromContext(ctx context.Context) Span {
 	return trace.SpanFromContext(ctx)
 }
 
+// ContextWithSpan returns a child context from ctx with span embedded
+func ContextWithSpan(ctx context.Context, span Span) context.Context {
+	return trace.ContextWithSpan(ctx, span)
+}
+
 // NewSpanFromContext creates a new span from the given context.
 // If the context already has a span attached to it, the new span will be a child of the existing span.
 // If the context does not have a span attached to it, the new span will be a root span.
diff --git a/trace/span_test.go b/trace/span_test.go
index caf97e7..c981b5c 100644
--- a/trace/span_test.go
+++ b/trace/span_test.go
@@ -51,6 +51,20 @@ func TestSpanFromContext(t *testing.T) {
 	}
 }
 
+func TestContextWithSpan(t *testing.T) {
+	provider, err := NewProvider(
+		WithContext(context.Background()),
+		WithConfig(&config.OpenTelemetry{Enabled: true}),
+	)
+	assert.NoError(t, err)
+	_, sampleSpan := provider.Tracer().Start(context.Background(), "sample span")
+
+	ctx := ContextWithSpan(context.Background(), sampleSpan)
+	if got := SpanFromContext(ctx); got != sampleSpan {
+		t.Errorf("ContextWithSpan() = %v, want %v", got, sampleSpan)
+	}
+}
+
 func TestNewSpanFromContext(t *testing.T) {
 	assert := assert.New(t)