test: New unit test for TraceDatabase function

scribd · Feb 22, 2023 · 69dfa41 · 69dfa41
1 parent 3c1371e
commit 69dfa41
Showing 1 changed file with 15 additions and 0 deletions.
diff --git a/pkg/instrumentation/database_test.go b/pkg/instrumentation/database_test.go
@@ -135,3 +135,18 @@ func TestInstrumentDatabase(t *testing.T) {
 		}
 	}
 }
+
+func TestTraceDatabase(t *testing.T) {
+	dbFile := path.Join(t.TempDir(), "test_db")
+	db, err := gorm.Open(sqlite.Open(dbFile))
+	if err != nil {
+		t.Fatalf("Failed to open DB: %s", err)
+	}
+
+	InstrumentDatabase(db, "test_app_name")
+	db = TraceDatabase(context.Background(), db)
+
+	if sp := db.Statement.Context.Value(parentSpanGormKey); sp == nil {
+		t.Error("Parent span not set on statement")
+	}
+}