Fix test_celery_setup error

[mediagoblin.git] / mediagoblin / tests / test_openid.py

diff --git a/mediagoblin/tests/test_openid.py b/mediagoblin/tests/test_openid.py
index 3aea79822db0d25f77bd4ff9e82a25dbaa29f007..71767032e892b54c8ef3cef431393f54269452d9 100644 (file)
--- a/mediagoblin/tests/test_openid.py
+++ b/mediagoblin/tests/test_openid.py
@@ -14,21 +14,26 @@
 # You should have received a copy of the GNU Affero General Public License
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
-import urlparse
 import pkg_resources
 import pytest
-import mock
+import six
+import six.moves.urllib.parse as urlparse
+try:
+    import mock
+except ImportError:
+    import unittest.mock as mock
 
 openid_consumer = pytest.importorskip(
     "openid.consumer.consumer")
 
 from mediagoblin import mg_globals
 from mediagoblin.db.base import Session
-from mediagoblin.db.models import User
+from mediagoblin.db.models import User, LocalUser
 from mediagoblin.plugins.openid.models import OpenIDUserURL
 from mediagoblin.tests.tools import get_app, fixture_add_user
 from mediagoblin.tools import template
 
+
 # App with plugin enabled
 @pytest.fixture()
 def openid_plugin_app(request):
@@ -177,7 +182,7 @@ class TestOpenIDPlugin(object):
 
             # Correct place?
             assert urlparse.urlsplit(res.location)[2] == '/u/chris/'
-            assert 'mediagoblin/user_pages/user.html' in template.TEMPLATE_TEST_CONTEXT
+            assert 'mediagoblin/user_pages/user_nonactive.html' in template.TEMPLATE_TEST_CONTEXT
 
             # No need to test if user is in logged in and verification email
             # awaits, since openid uses the register_user function which is
@@ -187,8 +192,9 @@ class TestOpenIDPlugin(object):
             openid_plugin_app.get('/auth/logout')
 
             # Get user and detach from session
-            test_user = mg_globals.database.User.query.filter_by(
-                username=u'chris').first()
+            test_user = mg_globals.database.LocalUser.query.filter(
+                LocalUser.username==u'chris'
+            ).first()
             Session.expunge(test_user)
 
             # Log back in
@@ -205,7 +211,7 @@ class TestOpenIDPlugin(object):
             # Make sure user is in the session
             context = template.TEMPLATE_TEST_CONTEXT['mediagoblin/root.html']
             session = context['request'].session
-            assert session['user_id'] == unicode(test_user.id)
+            assert session['user_id'] == six.text_type(test_user.id)
 
         _test_new_user()