From 55851ae23d60d1482f4499fb617dedd209db9f38 Mon Sep 17 00:00:00 2001
From: Sam Xie <sam@samxie.me>
Date: Sat, 26 Mar 2022 21:22:37 +0800
Subject: [PATCH] Add more tests

---
 sql.go      | 88 ++++++++++++++++++++++++++++-------------------------
 sql_test.go | 40 +++++++++++++++++++++---
 2 files changed, 81 insertions(+), 47 deletions(-)

diff --git a/sql.go b/sql.go
index 6b7032d..99b5be3 100644
--- a/sql.go
+++ b/sql.go
@@ -105,51 +105,55 @@ func RegisterDBStatsMetrics(db *sql.DB, dbSystem string, opts ...Option) error {
 	}, func(ctx context.Context) {
 		dbStats := db.Stats()
 
-		instruments.connectionMaxOpen.Observe(
-			ctx,
-			int64(dbStats.MaxOpenConnections),
-			cfg.Attributes...,
-		)
-
-		instruments.connectionOpen.Observe(
-			ctx,
-			int64(dbStats.InUse),
-			append(cfg.Attributes, connectionStatusKey.String("inuse"))...,
-		)
-		instruments.connectionOpen.Observe(
-			ctx,
-			int64(dbStats.Idle),
-			append(cfg.Attributes, connectionStatusKey.String("idle"))...,
-		)
-
-		instruments.connectionWaitTotal.Observe(
-			ctx,
-			dbStats.WaitCount,
-			cfg.Attributes...,
-		)
-		instruments.connectionWaitDurationTotal.Observe(
-			ctx,
-			float64(dbStats.WaitDuration.Nanoseconds())/1e6,
-			cfg.Attributes...,
-		)
-		instruments.connectionClosedMaxIdleTotal.Observe(
-			ctx,
-			dbStats.MaxIdleClosed,
-			cfg.Attributes...,
-		)
-		instruments.connectionClosedMaxIdleTimeTotal.Observe(
-			ctx,
-			dbStats.MaxIdleTimeClosed,
-			cfg.Attributes...,
-		)
-		instruments.connectionClosedMaxLifetimeTotal.Observe(
-			ctx,
-			dbStats.MaxLifetimeClosed,
-			cfg.Attributes...,
-		)
+		recordDBStatsMetrics(ctx, dbStats, instruments, cfg)
 	})
 	if err != nil {
 		return err
 	}
 	return nil
 }
+
+func recordDBStatsMetrics(ctx context.Context, dbStats sql.DBStats, instruments *dbStatsInstruments, cfg config) {
+	instruments.connectionMaxOpen.Observe(
+		ctx,
+		int64(dbStats.MaxOpenConnections),
+		cfg.Attributes...,
+	)
+
+	instruments.connectionOpen.Observe(
+		ctx,
+		int64(dbStats.InUse),
+		append(cfg.Attributes, connectionStatusKey.String("inuse"))...,
+	)
+	instruments.connectionOpen.Observe(
+		ctx,
+		int64(dbStats.Idle),
+		append(cfg.Attributes, connectionStatusKey.String("idle"))...,
+	)
+
+	instruments.connectionWaitTotal.Observe(
+		ctx,
+		dbStats.WaitCount,
+		cfg.Attributes...,
+	)
+	instruments.connectionWaitDurationTotal.Observe(
+		ctx,
+		float64(dbStats.WaitDuration.Nanoseconds())/1e6,
+		cfg.Attributes...,
+	)
+	instruments.connectionClosedMaxIdleTotal.Observe(
+		ctx,
+		dbStats.MaxIdleClosed,
+		cfg.Attributes...,
+	)
+	instruments.connectionClosedMaxIdleTimeTotal.Observe(
+		ctx,
+		dbStats.MaxIdleTimeClosed,
+		cfg.Attributes...,
+	)
+	instruments.connectionClosedMaxLifetimeTotal.Observe(
+		ctx,
+		dbStats.MaxLifetimeClosed,
+		cfg.Attributes...,
+	)
+}
diff --git a/sql_test.go b/sql_test.go
index 6800e14..20cd516 100644
--- a/sql_test.go
+++ b/sql_test.go
@@ -15,27 +15,37 @@
 package otelsql
 
 import (
+	"context"
 	"database/sql"
+	"fmt"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"go.opentelemetry.io/otel/attribute"
+	"go.opentelemetry.io/otel/metric/nonrecording"
 	semconv "go.opentelemetry.io/otel/semconv/v1.4.0"
 )
 
+var driverName string
+
 func init() {
 	sql.Register("test-driver", newMockDriver(false))
 	maxDriverSlot = 1
-}
 
-func TestRegister(t *testing.T) {
-	driverName, err := Register("test-driver", "test-db",
+	var err error
+	driverName, err = Register("test-driver", "test-db",
 		WithAttributes(attribute.String("foo", "bar")),
 	)
-	require.NoError(t, err)
-	assert.Equal(t, "test-driver-otelsql-0", driverName)
+	if err != nil {
+		panic(err)
+	}
+	if driverName != "test-driver-otelsql-0" {
+		panic(fmt.Sprintf("expect driver name: test-driver-otelsql-0, got %s", driverName))
+	}
+}
 
+func TestRegister(t *testing.T) {
 	// Expected driver
 	db, err := sql.Open(driverName, "")
 	require.NoError(t, err)
@@ -66,3 +76,23 @@ func TestWrapDriver(t *testing.T) {
 		attribute.String("foo", "bar"),
 	}, otelDriver.cfg.Attributes)
 }
+
+func TestRegisterDBStatsMetrics(t *testing.T) {
+	db, err := sql.Open(driverName, "")
+	require.NoError(t, err)
+
+	err = RegisterDBStatsMetrics(db, "test-db")
+	assert.NoError(t, err)
+}
+
+func TestRecordDBStatsMetricsNoPanic(t *testing.T) {
+	db, err := sql.Open(driverName, "")
+	require.NoError(t, err)
+
+	instruments, err := newDBStatsInstruments(nonrecording.NewNoopMeterProvider().Meter("test"))
+	require.NoError(t, err)
+
+	cfg := newConfig("db")
+
+	recordDBStatsMetrics(context.Background(), db.Stats(), instruments, cfg)
+}