Skip to content

Commit 53abbe4

Browse files
XSAMXSAM
authored
Add more tests for legacy Conn and Stmt (#159)
* Add tests for legacy Conn * Add tests for legacy Stmt * Add missing license
1 parent baa29c0 commit 53abbe4

File tree

4 files changed

+512
-188
lines changed

4 files changed

+512
-188
lines changed

‎conn_test.go

+175-112
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,19 @@ import (
2727
"go.opentelemetry.io/otel/trace"
2828
)
2929

30+
type MockConn interface {
31+
driver.Conn
32+
PrepareContextCount() int
33+
PrepareContextCtx() context.Context
34+
PrepareContextQuery() string
35+
36+
BeginTxCtx() context.Context
37+
BeginTxCount() int
38+
}
39+
40+
var _ MockConn = (*mockConn)(nil)
41+
var _ driver.Conn = (*mockConn)(nil)
42+
3043
type mockConn struct {
3144
shouldError bool
3245

@@ -52,6 +65,26 @@ type mockConn struct {
5265
pingCtx context.Context
5366
}
5467

68+
func (m *mockConn) BeginTxCtx() context.Context {
69+
return m.beginTxCtx
70+
}
71+
72+
func (m *mockConn) BeginTxCount() int {
73+
return m.beginTxCount
74+
}
75+
76+
func (m *mockConn) PrepareContextCount() int {
77+
return m.prepareContextCount
78+
}
79+
80+
func (m *mockConn) PrepareContextCtx() context.Context {
81+
return m.prepareContextCtx
82+
}
83+
84+
func (m *mockConn) PrepareContextQuery() string {
85+
return m.prepareContextQuery
86+
}
87+
5588
func newMockConn(shouldError bool) *mockConn {
5689
return &mockConn{shouldError: shouldError}
5790
}
@@ -406,90 +439,105 @@ func TestOtConn_PrepareContext(t *testing.T) {
406439
query := "query"
407440
expectedAttrs := []attribute.KeyValue{semconv.DBStatementKey.String(query)}
408441

409-
for _, omitConnPrepare := range []bool{true, false} {
442+
for _, legacy := range []bool{true, false} {
410443
var testname string
411-
if omitConnPrepare {
412-
testname = "OmitConnPrepare"
444+
if legacy {
445+
testname = "Legacy"
413446
}
414447

415448
t.Run(testname, func(t *testing.T) {
416-
testCases := []struct {
417-
name string
418-
error bool
419-
noParentSpan bool
420-
disableQuery bool
421-
attrs []attribute.KeyValue
422-
attributesGetter AttributesGetter
423-
}{
424-
{
425-
name: "no error",
426-
attrs: expectedAttrs,
427-
},
428-
{
429-
name: "no query db.statement",
430-
disableQuery: true,
431-
},
432-
{
433-
name: "with error",
434-
error: true,
435-
attrs: expectedAttrs,
436-
},
437-
{
438-
name: "no parent span",
439-
noParentSpan: true,
440-
attrs: expectedAttrs,
441-
},
442-
{
443-
name: "with attribute getter",
444-
attributesGetter: getDummyAttributesGetter(),
445-
attrs: expectedAttrs,
446-
},
447-
}
448-
449-
for _, tc := range testCases {
450-
t.Run(tc.name, func(t *testing.T) {
451-
// Prepare traces
452-
ctx, sr, tracer, dummySpan := prepareTraces(tc.noParentSpan)
453-
454-
// New conn
455-
cfg := newMockConfig(t, tracer)
456-
cfg.SpanOptions.DisableQuery = tc.disableQuery
457-
cfg.SpanOptions.OmitConnPrepare = omitConnPrepare
458-
cfg.AttributesGetter = tc.attributesGetter
459-
mc := newMockConn(tc.error)
460-
otelConn := newConn(mc, cfg)
449+
for _, omitConnPrepare := range []bool{true, false} {
450+
var testname string
451+
if omitConnPrepare {
452+
testname = "OmitConnPrepare"
453+
}
461454

462-
stmt, err := otelConn.PrepareContext(ctx, query)
463-
if tc.error {
464-
require.Error(t, err)
465-
} else {
466-
require.NoError(t, err)
455+
t.Run(testname, func(t *testing.T) {
456+
testCases := []struct {
457+
name string
458+
error bool
459+
noParentSpan bool
460+
disableQuery bool
461+
attrs []attribute.KeyValue
462+
attributesGetter AttributesGetter
463+
}{
464+
{
465+
name: "no error",
466+
attrs: expectedAttrs,
467+
},
468+
{
469+
name: "no query db.statement",
470+
disableQuery: true,
471+
},
472+
{
473+
name: "with error",
474+
error: true,
475+
attrs: expectedAttrs,
476+
},
477+
{
478+
name: "no parent span",
479+
noParentSpan: true,
480+
attrs: expectedAttrs,
481+
},
482+
{
483+
name: "with attribute getter",
484+
attributesGetter: getDummyAttributesGetter(),
485+
attrs: expectedAttrs,
486+
},
467487
}
468488

469-
spanList := sr.Ended()
470-
expectedSpanCount := getExpectedSpanCount(tc.noParentSpan, omitConnPrepare)
471-
// One dummy span and one span created in PrepareContext
472-
require.Equal(t, expectedSpanCount, len(spanList))
473-
474-
assertSpanList(t, spanList, spanAssertionParameter{
475-
parentSpan: dummySpan,
476-
error: tc.error,
477-
expectedAttributes: append(cfg.Attributes, tc.attrs...),
478-
method: MethodConnPrepare,
479-
noParentSpan: tc.noParentSpan,
480-
ctx: mc.prepareContextCtx,
481-
omitSpan: omitConnPrepare,
482-
attributesGetter: tc.attributesGetter,
483-
query: query,
484-
})
485-
486-
assert.Equal(t, 1, mc.prepareContextCount)
487-
assert.Equal(t, "query", mc.prepareContextQuery)
488-
489-
if !tc.error {
490-
otelStmt, ok := stmt.(*otStmt)
491-
require.True(t, ok)
492-
assert.Equal(t, "query", otelStmt.query)
489+
for _, tc := range testCases {
490+
t.Run(tc.name, func(t *testing.T) {
491+
// Prepare traces
492+
ctx, sr, tracer, dummySpan := prepareTraces(tc.noParentSpan)
493+
494+
// New conn
495+
cfg := newMockConfig(t, tracer)
496+
cfg.SpanOptions.DisableQuery = tc.disableQuery
497+
cfg.SpanOptions.OmitConnPrepare = omitConnPrepare
498+
cfg.AttributesGetter = tc.attributesGetter
499+
500+
var mc MockConn
501+
if legacy {
502+
mc = newMockLegacyConn(tc.error)
503+
} else {
504+
mc = newMockConn(tc.error)
505+
}
506+
otelConn := newConn(mc, cfg)
507+
508+
stmt, err := otelConn.PrepareContext(ctx, query)
509+
if tc.error {
510+
require.Error(t, err)
511+
} else {
512+
require.NoError(t, err)
513+
}
514+
515+
spanList := sr.Ended()
516+
expectedSpanCount := getExpectedSpanCount(tc.noParentSpan, omitConnPrepare)
517+
// One dummy span and one span created in PrepareContext
518+
require.Equal(t, expectedSpanCount, len(spanList))
519+
520+
assertSpanList(t, spanList, spanAssertionParameter{
521+
parentSpan: dummySpan,
522+
error: tc.error,
523+
expectedAttributes: append(cfg.Attributes, tc.attrs...),
524+
method: MethodConnPrepare,
525+
noParentSpan: tc.noParentSpan,
526+
ctx: mc.PrepareContextCtx(),
527+
omitSpan: omitConnPrepare,
528+
attributesGetter: tc.attributesGetter,
529+
query: query,
530+
})
531+
532+
assert.Equal(t, 1, mc.PrepareContextCount())
533+
assert.Equal(t, "query", mc.PrepareContextQuery())
534+
535+
if !tc.error {
536+
otelStmt, ok := stmt.(*otStmt)
537+
require.True(t, ok)
538+
assert.Equal(t, "query", otelStmt.query)
539+
}
540+
})
493541
}
494542
})
495543
}
@@ -521,48 +569,63 @@ func TestOtConn_BeginTx(t *testing.T) {
521569
},
522570
}
523571

524-
for _, tc := range testCases {
525-
t.Run(tc.name, func(t *testing.T) {
526-
// Prepare traces
527-
ctx, sr, tracer, dummySpan := prepareTraces(tc.noParentSpan)
572+
for _, legacy := range []bool{true, false} {
573+
var testname string
574+
if legacy {
575+
testname = "Legacy"
576+
}
528577

529-
// New conn
530-
cfg := newMockConfig(t, tracer)
531-
cfg.AttributesGetter = tc.attributesGetter
532-
mc := newMockConn(tc.error)
533-
otelConn := newConn(mc, cfg)
578+
t.Run(testname, func(t *testing.T) {
579+
for _, tc := range testCases {
580+
t.Run(tc.name, func(t *testing.T) {
581+
// Prepare traces
582+
ctx, sr, tracer, dummySpan := prepareTraces(tc.noParentSpan)
534583

535-
tx, err := otelConn.BeginTx(ctx, driver.TxOptions{})
536-
if tc.error {
537-
require.Error(t, err)
538-
} else {
539-
require.NoError(t, err)
540-
}
584+
// New conn
585+
cfg := newMockConfig(t, tracer)
586+
cfg.AttributesGetter = tc.attributesGetter
541587

542-
spanList := sr.Ended()
543-
expectedSpanCount := getExpectedSpanCount(tc.noParentSpan, false)
544-
// One dummy span and one span created in BeginTx
545-
require.Equal(t, expectedSpanCount, len(spanList))
588+
var mc MockConn
589+
if legacy {
590+
mc = newMockLegacyConn(tc.error)
591+
} else {
592+
mc = newMockConn(tc.error)
593+
}
594+
otelConn := newConn(mc, cfg)
546595

547-
assertSpanList(t, spanList, spanAssertionParameter{
548-
parentSpan: dummySpan,
549-
error: tc.error,
550-
expectedAttributes: cfg.Attributes,
551-
method: MethodConnBeginTx,
552-
noParentSpan: tc.noParentSpan,
553-
ctx: mc.beginTxCtx,
554-
attributesGetter: tc.attributesGetter,
555-
})
596+
tx, err := otelConn.BeginTx(ctx, driver.TxOptions{})
597+
if tc.error {
598+
require.Error(t, err)
599+
} else {
600+
require.NoError(t, err)
601+
}
602+
603+
spanList := sr.Ended()
604+
expectedSpanCount := getExpectedSpanCount(tc.noParentSpan, false)
605+
// One dummy span and one span created in BeginTx
606+
require.Equal(t, expectedSpanCount, len(spanList))
556607

557-
assert.Equal(t, 1, mc.beginTxCount)
608+
assertSpanList(t, spanList, spanAssertionParameter{
609+
parentSpan: dummySpan,
610+
error: tc.error,
611+
expectedAttributes: cfg.Attributes,
612+
method: MethodConnBeginTx,
613+
noParentSpan: tc.noParentSpan,
614+
ctx: mc.BeginTxCtx(),
615+
attributesGetter: tc.attributesGetter,
616+
})
558617

559-
if !tc.error {
560-
otelTx, ok := tx.(*otTx)
561-
require.True(t, ok)
618+
assert.Equal(t, 1, mc.BeginTxCount())
562619

563-
if dummySpan != nil {
564-
assert.Equal(t, dummySpan.SpanContext(), trace.SpanContextFromContext(otelTx.ctx))
565-
}
620+
if !tc.error {
621+
otelTx, ok := tx.(*otTx)
622+
require.True(t, ok)
623+
624+
if dummySpan != nil {
625+
assert.Equal(t, dummySpan.SpanContext(), trace.SpanContextFromContext(otelTx.ctx))
626+
}
627+
}
628+
})
566629
}
567630
})
568631
}

0 commit comments

Comments
 (0)