1 files changed, 16 insertions, 8 deletions
diff --git a/mediagoblin/tests/test_persona.py b/mediagoblin/tests/test_persona.py
index a1cd30eb..437cb7a1 100644
--- a/mediagoblin/tests/test_persona.py
+++ b/mediagoblin/tests/test_persona.py
@@ -13,16 +13,22 @@
 #
 # 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
+try:
+    import mock
+except ImportError:
+    import unittest.mock as mock
+
+import six.moves.urllib.parse as urlparse
 
 pytest.importorskip("requests")
 
 from mediagoblin import mg_globals
 from mediagoblin.db.base import Session
-from mediagoblin.db.models import Privilege
+from mediagoblin.db.models import Privilege, LocalUser
 from mediagoblin.tests.tools import get_app
 from mediagoblin.tools import template
 
@@ -111,14 +117,16 @@ class TestPersonaPlugin(object):
             persona_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()
             active_privilege = Privilege.query.filter(
                 Privilege.privilege_name==u'active').first()
             test_user.all_privileges.append(active_privilege)
             test_user.save()
-            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)
 
             # Add another user for _test_edit_persona
@@ -140,7 +148,7 @@ class TestPersonaPlugin(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_registration()