From ce3ead6b5fced43cd05c794e0847b50be285206e Mon Sep 17 00:00:00 2001
From: Daniel Pehrson <dpehrson@qurve.com>
Date: Fri, 21 Mar 2014 10:27:23 -0400
Subject: [PATCH] Ensure registration controller block yields happen on failure
 in addition to success and closes #2936.

Now with 100% more unit tests.
---
 .../devise/registrations_controller.rb        | 10 +++---
 .../custom_registrations_controller_test.rb   | 35 +++++++++++++++++++
 .../custom/registrations_controller.rb        | 21 +++++++++++
 test/rails_app/config/routes.rb               |  3 ++
 4 files changed, 65 insertions(+), 4 deletions(-)
 create mode 100644 test/controllers/custom_registrations_controller_test.rb
 create mode 100644 test/rails_app/app/controllers/custom/registrations_controller.rb

diff --git a/app/controllers/devise/registrations_controller.rb b/app/controllers/devise/registrations_controller.rb
index 07f43aa7..eab982e6 100644
--- a/app/controllers/devise/registrations_controller.rb
+++ b/app/controllers/devise/registrations_controller.rb
@@ -12,8 +12,9 @@ class Devise::RegistrationsController < DeviseController
   def create
     build_resource(sign_up_params)
 
-    if resource.save
-      yield resource if block_given?
+    resource_saved = resource.save
+    yield resource if block_given?
+    if resource_saved
       if resource.active_for_authentication?
         set_flash_message :notice, :signed_up if is_flashing_format?
         sign_up(resource_name, resource)
@@ -41,8 +42,9 @@ class Devise::RegistrationsController < DeviseController
     self.resource = resource_class.to_adapter.get!(send(:"current_#{resource_name}").to_key)
     prev_unconfirmed_email = resource.unconfirmed_email if resource.respond_to?(:unconfirmed_email)
 
-    if update_resource(resource, account_update_params)
-      yield resource if block_given?
+    resource_updated = update_resource(resource, account_update_params)
+    yield resource if block_given?
+    if resource_updated
       if is_flashing_format?
         flash_key = update_needs_confirmation?(resource, prev_unconfirmed_email) ?
           :update_needs_confirmation : :updated
diff --git a/test/controllers/custom_registrations_controller_test.rb b/test/controllers/custom_registrations_controller_test.rb
new file mode 100644
index 00000000..e56fddc7
--- /dev/null
+++ b/test/controllers/custom_registrations_controller_test.rb
@@ -0,0 +1,35 @@
+require 'test_helper'
+
+class CustomRegistrationsControllerTest < ActionController::TestCase
+  tests Custom::RegistrationsController
+
+  include Devise::TestHelpers
+
+  setup do
+    request.env["devise.mapping"] = Devise.mappings[:user]
+    @password = 'password'
+    @user = create_user(password: @password, password_confirmation: @password).tap(&:confirm!)
+  end
+
+  test "yield resource to block on create success" do
+    post :create, {user: {:email => "user@example.org", :password => "password", :password_confirmation => "password"}}
+    assert @controller.create_block_called?, "create failed to yield resource to provided block"
+  end
+
+  test "yield resource to block on create failure" do
+    post :create, {user: {}}
+    assert @controller.create_block_called?, "create failed to yield resource to provided block"
+  end
+
+  test "yield resource to block on update success" do
+    sign_in @user
+    put :update, {user: {current_password: @password}}
+    assert @controller.update_block_called?, "update failed to yield resource to provided block"
+  end
+
+  test "yield resource to block on update failure" do
+    sign_in @user
+    put :update, {user: {}}
+    assert @controller.update_block_called?, "update failed to yield resource to provided block"
+  end
+end
diff --git a/test/rails_app/app/controllers/custom/registrations_controller.rb b/test/rails_app/app/controllers/custom/registrations_controller.rb
new file mode 100644
index 00000000..9f1699c8
--- /dev/null
+++ b/test/rails_app/app/controllers/custom/registrations_controller.rb
@@ -0,0 +1,21 @@
+class Custom::RegistrationsController < Devise::RegistrationsController
+  def create
+    super do |resource|
+      @create_block_called = true
+    end
+  end
+
+  def update
+    super do |resource|
+      @update_block_called = true
+    end
+  end
+
+  def create_block_called?
+    @create_block_called == true
+  end
+
+  def update_block_called?
+    @update_block_called == true
+  end
+end
diff --git a/test/rails_app/config/routes.rb b/test/rails_app/config/routes.rb
index 16f60751..1ab9c290 100644
--- a/test/rails_app/config/routes.rb
+++ b/test/rails_app/config/routes.rb
@@ -26,6 +26,9 @@ Rails.application.routes.draw do
 
   get "/sign_in", to: "devise/sessions#new"
 
+  # Routes for custom controller testing
+  devise_for :user, only: [:registrations], controllers: { registrations: "custom/registrations" }, as: :custom, path: :custom
+
   # Admin scope
   devise_for :admin, path: "admin_area", controllers: { sessions: :"admins/sessions" }, skip: :passwords