updated sql instrumentation unit test with opt in changes

instana · Mar 25, 2024 · 0fba787 · 0fba787
1 parent 1f075b9
commit 0fba787
Show file tree

Hide file tree

Showing 3 changed files with 207 additions and 55 deletions.
diff --git a/instrumentation_sql_go1.10_test.go b/instrumentation_sql_go1.10_test.go
@@ -18,8 +18,7 @@ import (
 )
 
 func TestWrapSQLConnector_Exec(t *testing.T) {
-	beforeTest()
-	defer afterTest()
+
 	recorder := instana.NewTestRecorder()
 	s := instana.NewSensorWithTracer(instana.NewTracerWithEverything(&instana.Options{
 		Service:     "go-sensor-test",
@@ -29,7 +28,13 @@ func TestWrapSQLConnector_Exec(t *testing.T) {
 
 	db := sql.OpenDB(instana.WrapSQLConnector(s, "connection string", sqlConnector{}))
 
-	res, err := db.Exec("TEST QUERY")
+	pSpan := s.Tracer().StartSpan("parent-span")
+	ctx := context.Background()
+	if pSpan != nil {
+		ctx = instana.ContextWithSpan(ctx, pSpan)
+	}
+
+	res, err := db.ExecContext(ctx, "TEST QUERY")
 	require.NoError(t, err)
 
 	lastID, err := res.LastInsertId()
@@ -62,20 +67,25 @@ func TestWrapSQLConnector_Exec(t *testing.T) {
 }
 
 func TestWrapSQLConnector_Exec_Error(t *testing.T) {
-	beforeTest()
-	defer afterTest()
+
 	recorder := instana.NewTestRecorder()
 	s := instana.NewSensorWithTracer(instana.NewTracerWithEverything(&instana.Options{
 		Service:     "go-sensor-test",
 		AgentClient: alwaysReadyClient{},
 	}, recorder))
 	defer instana.ShutdownSensor()
 
+	pSpan := s.Tracer().StartSpan("parent-span")
+	ctx := context.Background()
+	if pSpan != nil {
+		ctx = instana.ContextWithSpan(ctx, pSpan)
+	}
+
 	db := sql.OpenDB(instana.WrapSQLConnector(s, "connection string", sqlConnector{
 		Error: errors.New("something went wrong"),
 	}))
 
-	_, err := db.Exec("TEST QUERY")
+	_, err := db.ExecContext(ctx, "TEST QUERY")
 	assert.Error(t, err)
 
 	spans := recorder.GetQueuedSpans()
@@ -105,18 +115,23 @@ func TestWrapSQLConnector_Exec_Error(t *testing.T) {
 }
 
 func TestWrapSQLConnector_Query(t *testing.T) {
-	beforeTest()
-	defer afterTest()
+
 	recorder := instana.NewTestRecorder()
 	s := instana.NewSensorWithTracer(instana.NewTracerWithEverything(&instana.Options{
 		Service:     "go-sensor-test",
 		AgentClient: alwaysReadyClient{},
 	}, recorder))
 	defer instana.ShutdownSensor()
 
+	pSpan := s.Tracer().StartSpan("parent-span")
+	ctx := context.Background()
+	if pSpan != nil {
+		ctx = instana.ContextWithSpan(ctx, pSpan)
+	}
+
 	db := sql.OpenDB(instana.WrapSQLConnector(s, "connection string", sqlConnector{}))
 
-	res, err := db.Query("TEST QUERY")
+	res, err := db.QueryContext(ctx, "TEST QUERY")
 	require.NoError(t, err)
 
 	cols, err := res.Columns()
@@ -149,8 +164,7 @@ func TestWrapSQLConnector_Query(t *testing.T) {
 }
 
 func TestWrapSQLConnector_Query_Error(t *testing.T) {
-	beforeTest()
-	defer afterTest()
+
 	recorder := instana.NewTestRecorder()
 	s := instana.NewSensorWithTracer(instana.NewTracerWithEverything(&instana.Options{
 		Service:     "go-sensor-test",
@@ -163,7 +177,13 @@ func TestWrapSQLConnector_Query_Error(t *testing.T) {
 		Error: dbErr,
 	}))
 
-	_, err := db.Query("TEST QUERY")
+	pSpan := s.Tracer().StartSpan("parent-span")
+	ctx := context.Background()
+	if pSpan != nil {
+		ctx = instana.ContextWithSpan(ctx, pSpan)
+	}
+
+	_, err := db.QueryContext(ctx, "TEST QUERY")
 	assert.Error(t, err)
 
 	spans := recorder.GetQueuedSpans()