refactor to improve database operability

jwt/revoke_test.go

@@ -1,6 +1,7 @@
 package jwt
 
 import (
+	"context"
 	"database/sql"
 	"testing"
 	"time"
@@ -20,13 +21,13 @@ func newGeneratorWithNoDB(t *testing.T) *TokenGenerator {
 		DB:                 nil,
 		DBType:             DatabaseType{Type: DatabasePostgreSQL, Version: "15"},
 		TableConfig:        DefaultTableConfig(),
-	})
+	}, nil)
 	require.NoError(t, err)
 
 	return gen
 }
 
-func newGeneratorWithMockDB(t *testing.T) (*TokenGenerator, sqlmock.Sqlmock, func()) {
+func newGeneratorWithMockDB(t *testing.T) (*TokenGenerator, *sql.DB, sqlmock.Sqlmock, func()) {
 	db, mock, err := sqlmock.New()
 	require.NoError(t, err)
 
@@ -34,6 +35,10 @@ func newGeneratorWithMockDB(t *testing.T) (*TokenGenerator, sqlmock.Sqlmock, fun
 	config.AutoCreate = false
 	config.EnableAutoCleanup = false
 
+	txGetter := func(ctx context.Context) (DBTransaction, error) {
+		return db.Begin()
+	}
+
 	gen, err := CreateGenerator(GeneratorConfig{
 		AccessExpireAfter:  15,
 		RefreshExpireAfter: 60,
@@ -43,10 +48,10 @@ func newGeneratorWithMockDB(t *testing.T) (*TokenGenerator, sqlmock.Sqlmock, fun
 		DB:                 db,
 		DBType:             DatabaseType{Type: DatabasePostgreSQL, Version: "15"},
 		TableConfig:        config,
-	})
+	}, txGetter)
 	require.NoError(t, err)
 
-	return gen, mock, func() { db.Close() }
+	return gen, db, mock, func() { db.Close() }
 }
 
 func TestNoDBFail(t *testing.T) {
@@ -73,7 +78,7 @@ func TestNoDBFail(t *testing.T) {
 }
 
 func TestRevokeAndCheckNotRevoked(t *testing.T) {
-	gen, mock, cleanup := newGeneratorWithMockDB(t)
+	gen, db, mock, cleanup := newGeneratorWithMockDB(t)
 	defer cleanup()
 
 	jti := uuid.New()
@@ -97,7 +102,7 @@ func TestRevokeAndCheckNotRevoked(t *testing.T) {
 		WillReturnRows(sqlmock.NewRows([]string{"1"}).AddRow(1))
 	mock.ExpectCommit()
 
-	tx, err := gen.db.Begin()
+	tx, err := db.Begin()
 	defer tx.Rollback()
 	require.NoError(t, err)