5 files changed, 23 insertions, 26 deletions
diff --git a/users/app/controllers/controller_extension/authentication.rb b/users/app/controllers/controller_extension/authentication.rb
index 50cf0d1..0408b77 100644
--- a/users/app/controllers/controller_extension/authentication.rb
+++ b/users/app/controllers/controller_extension/authentication.rb
@@ -8,7 +8,7 @@ module ControllerExtension::Authentication
   end
 
   def current_user
-    @current_user ||= env['warden'].user
+    @current_user ||= request.env['warden'].user
   end
 
   def logged_in?
diff --git a/users/test/functional/application_controller_test.rb b/users/test/functional/application_controller_test.rb
index 69bcb2f..857bae5 100644
--- a/users/test/functional/application_controller_test.rb
+++ b/users/test/functional/application_controller_test.rb
@@ -8,20 +8,19 @@ class ApplicationControllerTest < ActionController::TestCase
   end
 
   def test_authorize_redirect
-    stub_logged_out
     @controller.send(:authorize)
     assert_access_denied
   end
 
   def test_authorized
-    @user = stub_logged_in
+    login
     @controller.send(:authorize)
     assert_access_denied(false)
   end
 
   def test_authorize_admin
-    @user = stub_logged_in
-    @user.expects(:is_admin?).returns(false)
+    login
+    @current_user.expects(:is_admin?).returns(false)
     @controller.send(:authorize_admin)
     assert_access_denied
   end
diff --git a/users/test/functional/helper_methods_test.rb b/users/test/functional/helper_methods_test.rb
index c0eaf61..2b2375c 100644
--- a/users/test/functional/helper_methods_test.rb
+++ b/users/test/functional/helper_methods_test.rb
@@ -16,26 +16,23 @@ class HelperMethodsTest < ActionController::TestCase
     @controller
   end
 
-  def test_current_user_with_caching
-    @user = stub_logged_in
-    assert_equal @user, current_user
-    assert_equal @user, current_user # tests caching
+  def test_current_user
+    login
+    assert_equal @current_user, current_user
   end
 
   def test_logged_in
-    @user = stub_logged_in
+    login
     assert logged_in?
   end
 
   def test_logged_out
-    stub_logged_out
     assert !logged_in?
   end
 
   def test_admin
-    bool = stub
-    @user = stub_logged_in
-    @user.expects(:is_admin?).returns(bool)
+    login
+    @current_user.expects(:is_admin?).returns(bool = stub)
     assert_equal bool, admin?
   end
 
diff --git a/users/test/integration/api/account_flow_test.rb b/users/test/integration/api/account_flow_test.rb
index 5800d46..69e0599 100644
--- a/users/test/integration/api/account_flow_test.rb
+++ b/users/test/integration/api/account_flow_test.rb
@@ -4,7 +4,7 @@ class AccountFlowTest < ActionDispatch::IntegrationTest
 
   # this test wraps the api and implements the interface the ruby-srp client.
   def handshake(login, aa)
-    post "sessions", :login => login, 'A' => aa.to_s(16)
+    post "sessions", :login => login, 'A' => aa.to_s(16), :format => :json
     assert_response :success
     response = JSON.parse(@response.body)
     if response['errors']
@@ -15,7 +15,7 @@ class AccountFlowTest < ActionDispatch::IntegrationTest
   end
 
   def validate(m)
-    put "sessions/" + @login, :client_auth => m.to_s(16)
+    put "sessions/" + @login, :client_auth => m.to_s(16), :format => :json
     assert_response :success
     return JSON.parse(@response.body)
   end
diff --git a/users/test/support/auth_test_helper.rb b/users/test/support/auth_test_helper.rb
index 9412058..f211597 100644
--- a/users/test/support/auth_test_helper.rb
+++ b/users/test/support/auth_test_helper.rb
@@ -1,17 +1,18 @@
 module AuthTestHelper
+  extend ActiveSupport::Concern
 
-  def stub_logged_in
-    @user_id = stub
-    @user = stub
-    session[:user_id] = @user_id
-    User.expects(:find).once.with(@user_id).returns(@user)
-    return @user
+  # Controller will fetch current user from warden.
+  # Make it pick up our current_user
+  included do
+    setup do
+      request.env['warden'] ||= stub :user => nil
+    end
   end
 
-  def stub_logged_out
-    @user_id = stub
-    session[:user_id] = @user_id
-    User.expects(:find).once.with(@user_id).returns(nil)
+  def login(user = nil)
+    @current_user = user || stub
+    request.env['warden'] = stub :user => @current_user
+    return @current_user
   end
 
   def assert_access_denied(denied = true)