db: use `context` and go-mockgen for `TwoFactorsStore` (#7045)

1 files changed, 32 insertions, 38 deletions

diff --git a/internal/db/two_factors_test.go b/internal/db/two_factors_test.go
index c8412213..acd9a576 100644
--- a/internal/db/two_factors_test.go
+++ b/internal/db/two_factors_test.go
@@ -5,15 +5,17 @@
 package db
 
 import (
+	"context"
 	"testing"
 	"time"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 
 	"gogs.io/gogs/internal/errutil"
 )
 
-func Test_twoFactors(t *testing.T) {
+func TestTwoFactors(t *testing.T) {
 	if testing.Short() {
 		t.Skip()
 	}
@@ -29,16 +31,14 @@ func Test_twoFactors(t *testing.T) {
 		name string
 		test func(*testing.T, *twoFactors)
 	}{
-		{"Create", test_twoFactors_Create},
-		{"GetByUserID", test_twoFactors_GetByUserID},
-		{"IsUserEnabled", test_twoFactors_IsUserEnabled},
+		{"Create", twoFactorsCreate},
+		{"GetByUserID", twoFactorsGetByUserID},
+		{"IsUserEnabled", twoFactorsIsUserEnabled},
 	} {
 		t.Run(tc.name, func(t *testing.T) {
 			t.Cleanup(func() {
 				err := clearTables(t, db.DB, tables...)
-				if err != nil {
-					t.Fatal(err)
-				}
+				require.NoError(t, err)
 			})
 			tc.test(t, db)
 		})
@@ -48,55 +48,49 @@ func Test_twoFactors(t *testing.T) {
 	}
 }
 
-func test_twoFactors_Create(t *testing.T, db *twoFactors) {
+func twoFactorsCreate(t *testing.T, db *twoFactors) {
+	ctx := context.Background()
+
 	// Create a 2FA token
-	err := db.Create(1, "secure-key", "secure-secret")
-	if err != nil {
-		t.Fatal(err)
-	}
+	err := db.Create(ctx, 1, "secure-key", "secure-secret")
+	require.NoError(t, err)
 
 	// Get it back and check the Created field
-	tf, err := db.GetByUserID(1)
-	if err != nil {
-		t.Fatal(err)
-	}
+	tf, err := db.GetByUserID(ctx, 1)
+	require.NoError(t, err)
 	assert.Equal(t, db.NowFunc().Format(time.RFC3339), tf.Created.UTC().Format(time.RFC3339))
 
 	// Verify there are 10 recover codes generated
 	var count int64
 	err = db.Model(new(TwoFactorRecoveryCode)).Count(&count).Error
-	if err != nil {
-		t.Fatal(err)
-	}
+	require.NoError(t, err)
 	assert.Equal(t, int64(10), count)
 }
 
-func test_twoFactors_GetByUserID(t *testing.T, db *twoFactors) {
+func twoFactorsGetByUserID(t *testing.T, db *twoFactors) {
+	ctx := context.Background()
+
 	// Create a 2FA token for user 1
-	err := db.Create(1, "secure-key", "secure-secret")
-	if err != nil {
-		t.Fatal(err)
-	}
+	err := db.Create(ctx, 1, "secure-key", "secure-secret")
+	require.NoError(t, err)
 
 	// We should be able to get it back
-	_, err = db.GetByUserID(1)
-	if err != nil {
-		t.Fatal(err)
-	}
+	_, err = db.GetByUserID(ctx, 1)
+	require.NoError(t, err)
 
 	// Try to get a non-existent 2FA token
-	_, err = db.GetByUserID(2)
-	expErr := ErrTwoFactorNotFound{args: errutil.Args{"userID": int64(2)}}
-	assert.Equal(t, expErr, err)
+	_, err = db.GetByUserID(ctx, 2)
+	wantErr := ErrTwoFactorNotFound{args: errutil.Args{"userID": int64(2)}}
+	assert.Equal(t, wantErr, err)
 }
 
-func test_twoFactors_IsUserEnabled(t *testing.T, db *twoFactors) {
+func twoFactorsIsUserEnabled(t *testing.T, db *twoFactors) {
+	ctx := context.Background()
+
 	// Create a 2FA token for user 1
-	err := db.Create(1, "secure-key", "secure-secret")
-	if err != nil {
-		t.Fatal(err)
-	}
+	err := db.Create(ctx, 1, "secure-key", "secure-secret")
+	require.NoError(t, err)
 
-	assert.True(t, db.IsUserEnabled(1))
-	assert.False(t, db.IsUserEnabled(2))
+	assert.True(t, db.IsUserEnabled(ctx, 1))
+	assert.False(t, db.IsUserEnabled(ctx, 2))
 }