From 15b3f4d776cb06457ddd5679a31a09a2bfcb11cd Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Robert=20Paj=C4=85k?= <pellared@hotmail.com>
Date: Fri, 9 Feb 2024 09:01:32 +0100
Subject: [PATCH] Fix TestWithIDGenerator (#4894)

---
 sdk/trace/trace_test.go | 36 ++++++++++++++++++++----------------
 1 file changed, 20 insertions(+), 16 deletions(-)

diff --git a/sdk/trace/trace_test.go b/sdk/trace/trace_test.go
index e101e78f4..df8dab7d6 100644
--- a/sdk/trace/trace_test.go
+++ b/sdk/trace/trace_test.go
@@ -1954,25 +1954,29 @@ var _ IDGenerator = (*testIDGenerator)(nil)
 func TestWithIDGenerator(t *testing.T) {
 	const (
 		startTraceID = 1
-		startSpanID  = 1
-		numSpan      = 10
+		startSpanID  = 10
+		numSpan      = 5
 	)
 
-	gen := &testIDGenerator{traceID: startSpanID, spanID: startSpanID}
-
+	gen := &testIDGenerator{traceID: startTraceID, spanID: startSpanID}
+	te := NewTestExporter()
+	tp := NewTracerProvider(
+		WithSyncer(te),
+		WithIDGenerator(gen),
+	)
 	for i := 0; i < numSpan; i++ {
-		te := NewTestExporter()
-		tp := NewTracerProvider(
-			WithSyncer(te),
-			WithIDGenerator(gen),
-		)
-		span := startSpan(tp, "TestWithIDGenerator")
-		got, err := strconv.ParseUint(span.SpanContext().SpanID().String(), 16, 64)
-		require.NoError(t, err)
-		want := uint64(startSpanID + i)
-		assert.Equal(t, got, want)
-		_, err = endSpan(te, span)
-		require.NoError(t, err)
+		func() {
+			_, span := tp.Tracer(t.Name()).Start(context.Background(), strconv.Itoa(i))
+			defer span.End()
+
+			gotSpanID, err := strconv.ParseUint(span.SpanContext().SpanID().String(), 16, 64)
+			require.NoError(t, err)
+			assert.Equal(t, uint64(startSpanID+i), gotSpanID)
+
+			gotTraceID, err := strconv.ParseUint(span.SpanContext().TraceID().String(), 16, 64)
+			require.NoError(t, err)
+			assert.Equal(t, uint64(startTraceID+i), gotTraceID)
+		}()
 	}
 }