specs for forced two-factor authentication and grace period

simplified code and fixed stuffs

2 files changed, 11 insertions, 8 deletions

diff --git a/app/controllers/application_controller.rb b/app/controllers/application_controller.rb
index 978a269ca52..a945b38e35f 100644
--- a/app/controllers/application_controller.rb
+++ b/app/controllers/application_controller.rb
@@ -226,12 +226,7 @@ class ApplicationController < ActionController::Base
 
   def check_tfa_requirement
     if two_factor_authentication_required? && current_user && !current_user.two_factor_enabled && !skip_two_factor?
-      grace_period_started = current_user.otp_grace_period_started_at
-      grace_period_deadline = grace_period_started + two_factor_grace_period.hours
-
-      deadline_text = "until #{l(grace_period_deadline)}" unless two_factor_grace_period_expired?(grace_period_started)
-      redirect_to new_profile_two_factor_auth_path,
-                  alert: "You must configure Two-Factor Authentication in your account #{deadline_text}"
+      redirect_to new_profile_two_factor_auth_path
     end
   end
 
@@ -377,7 +372,8 @@ class ApplicationController < ActionController::Base
     current_application_settings.two_factor_grace_period
   end
 
-  def two_factor_grace_period_expired?(date)
+  def two_factor_grace_period_expired?
+    date = current_user.otp_grace_period_started_at
     date && (date + two_factor_grace_period.hours) < Time.current
   end
 
diff --git a/app/controllers/profiles/two_factor_auths_controller.rb b/app/controllers/profiles/two_factor_auths_controller.rb
index 49629e9894a..4f125eb7e05 100644
--- a/app/controllers/profiles/two_factor_auths_controller.rb
+++ b/app/controllers/profiles/two_factor_auths_controller.rb
@@ -10,6 +10,13 @@ class Profiles::TwoFactorAuthsController < Profiles::ApplicationController
     end
     current_user.save! if current_user.changed?
 
+    if two_factor_grace_period_expired?
+      flash.now[:alert] = 'You must configure Two-Factor Authentication in your account.'
+    else
+      grace_period_deadline = current_user.otp_grace_period_started_at + two_factor_grace_period.hours
+      flash.now[:alert] = "You must configure Two-Factor Authentication in your account until #{l(grace_period_deadline)}."
+    end
+
     @qr_code = build_qr_code
   end
 
@@ -40,7 +47,7 @@ class Profiles::TwoFactorAuthsController < Profiles::ApplicationController
   end
 
   def skip
-    if two_factor_grace_period_expired?(current_user.otp_grace_period_started_at)
+    if two_factor_grace_period_expired?
       redirect_to new_profile_two_factor_auth_path, alert: 'Cannot skip two factor authentication setup'
     else
       session[:skip_tfa] = current_user.otp_grace_period_started_at + two_factor_grace_period.hours