From f9255b7ccb34c23caa8a29e99b658cf4ac2c03df Mon Sep 17 00:00:00 2001
From: scambra <sergio@enpijama.es>
Date: Mon, 23 Apr 2012 13:26:55 +0200
Subject: [PATCH] Fix #203, keep invitation info on registering only for
 invitable resources

---
 .../controllers/registrations.rb              | 31 +++++++++++--------
 .../registrations_controller_test.rb          | 13 +++++++-
 test/rails_app/app/models/admin.rb            |  2 +-
 test/test_helper.rb                           |  3 ++
 4 files changed, 34 insertions(+), 15 deletions(-)

diff --git a/lib/devise_invitable/controllers/registrations.rb b/lib/devise_invitable/controllers/registrations.rb
index 9cc53241..29e09b47 100644
--- a/lib/devise_invitable/controllers/registrations.rb
+++ b/lib/devise_invitable/controllers/registrations.rb
@@ -1,35 +1,40 @@
 module DeviseInvitable::Controllers::Registrations
   def self.included(controller)
-    controller.send :around_filter, :destroy_if_previously_invited, :only => :create
+    controller.send :around_filter, :keep_invitation_info, :only => :create
   end
 
   protected
 
   def destroy_if_previously_invited
-    invitation_info = {}
+    @invitation_info = {}
 
     hash = params[resource_name]
     if hash && hash[:email]
       resource = resource_class.where(:email => hash[:email], :encrypted_password => '').first
       if resource
-        invitation_info[:invitation_sent_at] = resource[:invitation_sent_at]
-        invitation_info[:invited_by_id] = resource[:invited_by_id]
-        invitation_info[:invited_by_type] = resource[:invited_by_type]
+        @invitation_info[:invitation_sent_at] = resource[:invitation_sent_at]
+        @invitation_info[:invited_by_id] = resource[:invited_by_id]
+        @invitation_info[:invited_by_type] = resource[:invited_by_type]
         resource.destroy
       end
     end
-
-    # execute the action (create)
+  end
+  
+  def keep_invitation_info
+    resource_invitable = resource_class.devise_modules.include?(:invitable)
+    destroy_if_previously_invited if resource_invitable
     yield
-    # Note that the after_filter is executed at THIS position !
-
+    reset_invitation_info if resource_invitable
+  end
+  
+  def reset_invitation_info
     # Restore info about the last invitation (for later reference)
     # Reset the invitation_info only, if invited_by_id is still nil at this stage:
-    resource = resource_class.where(:email => hash[:email], :invited_by_id => nil).first
+    resource = resource_class.where(:email => params[resource_name][:email], :invited_by_id => nil).first
     if resource
-      resource[:invitation_sent_at] = invitation_info[:invitation_sent_at]
-      resource[:invited_by_id] = invitation_info[:invited_by_id]
-      resource[:invited_by_type] = invitation_info[:invited_by_type]
+      resource[:invitation_sent_at] = @invitation_info[:invitation_sent_at]
+      resource[:invited_by_id] = @invitation_info[:invited_by_id]
+      resource[:invited_by_type] = @invitation_info[:invited_by_type]
       resource.save!
     end
   end
diff --git a/test/functional/registrations_controller_test.rb b/test/functional/registrations_controller_test.rb
index ff35341f..0da409a9 100644
--- a/test/functional/registrations_controller_test.rb
+++ b/test/functional/registrations_controller_test.rb
@@ -1,4 +1,5 @@
 require 'test_helper'
+require 'model_tests_helper'
 
 class Devise::RegistrationsControllerTest < ActionController::TestCase
   def setup
@@ -27,7 +28,7 @@ def setup
     assert_blank @invitee.encrypted_password, "the password should be unset"
 
     # sign_up the invitee
-    post :create, :user => { :email => invitee_email, :password => "1password"}
+    post :create, :user => {:email => invitee_email, :password => "1password"}
 
     @invitee = User.where(:email => invitee_email).first
     assert_present @invitee.encrypted_password
@@ -36,4 +37,14 @@ def setup
     assert_present @invitee.invited_by_id
     assert_present @invitee.invited_by_type
   end
+  
+  test "not invitable resources can register" do
+    @request.env["devise.mapping"] = Devise.mappings[:admin]
+    invitee_email = "invitee@example.org"
+    
+    post :create, :admin => {:email => invitee_email, :password => "1password"}
+    
+    @invitee = Admin.where(:email => invitee_email).first
+    assert_present @invitee.encrypted_password
+  end
 end
diff --git a/test/rails_app/app/models/admin.rb b/test/rails_app/app/models/admin.rb
index b0ed5430..b094e89a 100644
--- a/test/rails_app/app/models/admin.rb
+++ b/test/rails_app/app/models/admin.rb
@@ -15,6 +15,6 @@ class Admin < PARENT_MODEL_CLASS
 
 
 
-  devise :database_authenticatable, :validatable
+  devise :database_authenticatable, :validatable, :registerable
   include DeviseInvitable::Inviter
 end
diff --git a/test/test_helper.rb b/test/test_helper.rb
index 68d5d038..d0312b68 100644
--- a/test/test_helper.rb
+++ b/test/test_helper.rb
@@ -18,3 +18,6 @@
 class ActionDispatch::IntegrationTest
   include Capybara::DSL
 end
+class ActionController::TestCase 
+  include Devise::TestHelpers 
+end
\ No newline at end of file