3 files changed, 57 insertions, 31 deletions

diff --git a/users/test/functional/users_controller_test.rb b/users/test/functional/users_controller_test.rb
index 44b6768..939d105 100644
--- a/users/test/functional/users_controller_test.rb
+++ b/users/test/functional/users_controller_test.rb
@@ -5,6 +5,7 @@ class UsersControllerTest < ActionController::TestCase
 
   test "should get new" do
     get :new
+
     assert_equal User, assigns(:user).class
     assert_response :success
   end
@@ -12,7 +13,9 @@ class UsersControllerTest < ActionController::TestCase
   test "should create new user" do
     user = stub_record User
     User.expects(:create).with(user.params).returns(user)
+
     post :create, :user => user.params, :format => :json
+
     assert_nil session[:user_id]
     assert_json_response user
     assert_response :success
@@ -24,70 +27,81 @@ class UsersControllerTest < ActionController::TestCase
     params.stringify_keys!
     assert !user.valid?
     User.expects(:create).with(params).returns(user)
+
     post :create, :user => params, :format => :json
+
     assert_json_error user.errors.messages
     assert_response 422
   end
 
   test "should get edit view" do
-    user = stub_record User
-    User.expects(:find_by_param).with(user.id.to_s).returns(user)
+    user = find_record User
+
     login user
     get :edit, :id => user.id
+
     assert_equal user, assigns[:user]
   end
 
   test "should process updated params" do
-    user = stub_record User
+    user = find_record User
     user.expects(:update_attributes).with(user.params).returns(true)
-    User.expects(:find_by_param).with(user.id.to_s).returns(user)
+
     login user
     put :update, :user => user.params, :id => user.id, :format => :json
+
     assert_equal user, assigns[:user]
     assert_response 204
     assert_equal " ", @response.body
   end
 
-  test "admin can edit user" do
-    user = stub_record User
+  test "admin can update user" do
+    user = find_record User
     user.expects(:update_attributes).with(user.params).returns(true)
-    User.expects(:find_by_param).with(user.id.to_s).returns(user)
+
     login :is_admin? => true
     put :update, :user => user.params, :id => user.id, :format => :json
+
     assert_equal user, assigns[:user]
     assert_response 204
     assert_equal " ", @response.body
   end
 
   test "admin can destroy user" do
-    login :is_admin? => true
-    user = stub_record User
+    user = find_record User
     user.expects(:destroy)
-    User.expects(:find_by_param).with(user.id).returns(user)
+
+    login :is_admin? => true
     delete :destroy, :id => user.id
+
     assert_response :redirect
     assert_redirected_to users_path
   end
 
   test "user can cancel account" do
-    login
-    @current_user.expects(:destroy)
-    User.expects(:find_by_param).with(@current_user.id).returns(@current_user)
+    user = find_record User
+    user.expects(:destroy)
+
+    login user
     delete :destroy, :id => @current_user.id
+
     assert_response :redirect
     assert_redirected_to login_path
   end
 
   test "non-admin can't destroy user" do
-    login
     user = stub_record User
+
+    login
     delete :destroy, :id => user.id
+
     assert_access_denied
   end
 
   test "admin can list users" do
     login :is_admin? => true
     get :index
+
     assert_response :success
     assert assigns(:users)
   end
@@ -95,12 +109,14 @@ class UsersControllerTest < ActionController::TestCase
   test "non-admin can't list users" do
     login
     get :index
+
     assert_access_denied
   end
 
   test "admin can autocomplete users" do
     login :is_admin? => true
     get :index, :format => :json
+
     assert_response :success
     assert assigns(:users)
   end
@@ -108,6 +124,7 @@ class UsersControllerTest < ActionController::TestCase
   test "admin can search users" do
     login :is_admin? => true
     get :index, :query => "a"
+
     assert_response :success
     assert assigns(:users)
   end
diff --git a/users/test/support/auth_test_helper.rb b/users/test/support/auth_test_helper.rb
index e0b673a..f3506ae 100644
--- a/users/test/support/auth_test_helper.rb
+++ b/users/test/support/auth_test_helper.rb
@@ -10,8 +10,8 @@ module AuthTestHelper
     end
   end
 
-  def login(user_or_method_hash = nil)
-    @current_user = stub_user(user_or_method_hash)
+  def login(user_or_method_hash = {})
+    @current_user = stub_record(User, user_or_method_hash)
     unless @current_user.respond_to? :is_admin?
       @current_user.stubs(:is_admin?).returns(false)
     end
@@ -28,19 +28,6 @@ module AuthTestHelper
     end
   end
 
-  protected
-
-  # Will create a stub user for logging in from either
-  # * a hash of methods to stub
-  # * a user record
-  # * nil -> create a user record stub
-  def stub_user(user_or_method_hash)
-    if user_or_method_hash.is_a?(Hash)
-      stub_record User, user_or_method_hash
-    else
-      user_or_method_hash || stub_record(User)
-    end
-  end
 end
 
 class ActionController::TestCase
diff --git a/users/test/support/stub_record_helper.rb b/users/test/support/stub_record_helper.rb
index ede21cf..2e1a533 100644
--- a/users/test/support/stub_record_helper.rb
+++ b/users/test/support/stub_record_helper.rb
@@ -1,19 +1,41 @@
 module StubRecordHelper
 
+  # Will expect find_by_param or find_by_id to be called on klass and
+  # return the record given.
+  # If no record is given but a hash or nil will create a stub based on
+  # that instead and returns the stub.
+  def find_record(klass, record_or_method_hash = {})
+    record = stub_record(klass, record_or_method_hash)
+    finder = klass.respond_to?(:find_by_param) ? :find_by_param : :find_by_id
+    klass.expects(finder).with(record.to_param).returns(record)
+    return record
+  end
+
   # Create a stub that has the usual functions of a database record.
   # It won't fail on rendering a form for example.
-  def stub_record(klass, params = {}, persisted = true)
+  #
+  # If the second parameter is a record we return the record itself.
+  # This way you can build functions that either take a record or a
+  # method hash to stub from. See find_record for an example.
+  def stub_record(klass, record_or_method_hash = {}, persisted = true)
+    if record_or_method_hash && !record_or_method_hash.is_a?(Hash)
+      return record_or_method_hash
+    end
+    stub record_params_for(klass, record_or_method_hash, persisted)
+  end
+
+  def record_params_for(klass, params = {}, persisted = true)
     if klass.respond_to?(:valid_attributes_hash)
       params.reverse_merge!(klass.valid_attributes_hash)
     end
     params[:params] = params.stringify_keys
     params.reverse_merge! :id => "A123",
+      :to_param => "A123",
       :class => klass,
       :to_key => ['123'],
       :to_json => %Q({"stub":"#{klass.name}"}),
       :new_record? => !persisted,
       :persisted? => persisted
-    stub params
   end
 
 end