3 files changed, 15 insertions, 13 deletions
diff --git a/internal/route/lfs/route_test.go b/internal/route/lfs/route_test.go
index 6202695b..ee668fc6 100644
--- a/internal/route/lfs/route_test.go
+++ b/internal/route/lfs/route_test.go
@@ -31,7 +31,7 @@ func Test_authenticate(t *testing.T) {
 		name                  string
 		header                http.Header
 		mockUsersStore        func() db.UsersStore
-		mockTwoFactorsStore   *db.MockTwoFactorsStore
+		mockTwoFactorsStore   func() db.TwoFactorsStore
 		mockAccessTokensStore func() db.AccessTokensStore
 		expStatusCode         int
 		expHeader             http.Header
@@ -56,10 +56,10 @@ func Test_authenticate(t *testing.T) {
 				mock.AuthenticateFunc.SetDefaultReturn(&db.User{}, nil)
 				return mock
 			},
-			mockTwoFactorsStore: &db.MockTwoFactorsStore{
-				MockIsUserEnabled: func(userID int64) bool {
-					return true
-				},
+			mockTwoFactorsStore: func() db.TwoFactorsStore {
+				mock := db.NewMockTwoFactorsStore()
+				mock.IsUserEnabledFunc.SetDefaultReturn(true)
+				return mock
 			},
 			expStatusCode: http.StatusBadRequest,
 			expHeader:     http.Header{},
@@ -98,10 +98,10 @@ func Test_authenticate(t *testing.T) {
 				mock.AuthenticateFunc.SetDefaultReturn(&db.User{ID: 1, Name: "unknwon"}, nil)
 				return mock
 			},
-			mockTwoFactorsStore: &db.MockTwoFactorsStore{
-				MockIsUserEnabled: func(userID int64) bool {
-					return false
-				},
+			mockTwoFactorsStore: func() db.TwoFactorsStore {
+				mock := db.NewMockTwoFactorsStore()
+				mock.IsUserEnabledFunc.SetDefaultReturn(false)
+				return mock
 			},
 			expStatusCode: http.StatusOK,
 			expHeader:     http.Header{},
@@ -133,7 +133,9 @@ func Test_authenticate(t *testing.T) {
 			if test.mockUsersStore != nil {
 				db.SetMockUsersStore(t, test.mockUsersStore())
 			}
-			db.SetMockTwoFactorsStore(t, test.mockTwoFactorsStore)
+			if test.mockTwoFactorsStore != nil {
+				db.SetMockTwoFactorsStore(t, test.mockTwoFactorsStore())
+			}
 			if test.mockAccessTokensStore != nil {
 				db.SetMockAccessTokensStore(t, test.mockAccessTokensStore())
 			}
diff --git a/internal/route/user/auth.go b/internal/route/user/auth.go
index b3b785c2..cae46853 100644
--- a/internal/route/user/auth.go
+++ b/internal/route/user/auth.go
@@ -210,7 +210,7 @@ func LoginTwoFactorPost(c *context.Context) {
 		return
 	}
 
-	t, err := db.TwoFactors.GetByUserID(userID)
+	t, err := db.TwoFactors.GetByUserID(c.Req.Context(), userID)
 	if err != nil {
 		c.Error(err, "get two factor by user ID")
 		return
diff --git a/internal/route/user/setting.go b/internal/route/user/setting.go
index c28a1747..66e3ace2 100644
--- a/internal/route/user/setting.go
+++ b/internal/route/user/setting.go
@@ -381,7 +381,7 @@ func SettingsSecurity(c *context.Context) {
 	c.Title("settings.security")
 	c.PageIs("SettingsSecurity")
 
-	t, err := db.TwoFactors.GetByUserID(c.UserID())
+	t, err := db.TwoFactors.GetByUserID(c.Req.Context(), c.UserID())
 	if err != nil && !db.IsErrTwoFactorNotFound(err) {
 		c.Errorf(err, "get two factor by user ID")
 		return
@@ -449,7 +449,7 @@ func SettingsTwoFactorEnablePost(c *context.Context) {
 		return
 	}
 
-	if err := db.TwoFactors.Create(c.UserID(), conf.Security.SecretKey, secret); err != nil {
+	if err := db.TwoFactors.Create(c.Req.Context(), c.UserID(), conf.Security.SecretKey, secret); err != nil {
 		c.Flash.Error(c.Tr("settings.two_factor_enable_error", err))
 		c.RedirectSubpath("/user/settings/security/two_factor_enable")
 		return