1 files changed, 4 insertions, 3 deletions
diff --git a/tests/regressiontests/multiple_database/tests.py b/tests/regressiontests/multiple_database/tests.py
index 595c5edb3b..74a5f2f550 100644
--- a/tests/regressiontests/multiple_database/tests.py
+++ b/tests/regressiontests/multiple_database/tests.py
@@ -2,7 +2,7 @@ from __future__ import absolute_import, unicode_literals
 
 import datetime
 import pickle
-from StringIO import StringIO
+from operator import attrgetter
 
 from django.conf import settings
 from django.contrib.auth.models import User
@@ -11,6 +11,7 @@ from django.core import management
 from django.db import connections, router, DEFAULT_DB_ALIAS
 from django.db.models import signals
 from django.test import TestCase
+from django.utils.six import StringIO
 
 from .models import Book, Person, Pet, Review, UserProfile
 
@@ -873,10 +874,10 @@ class QueryTestCase(TestCase):
         dive = Book.objects.using('other').create(title="Dive into Python",
             published=datetime.date(2009, 5, 4))
         val = Book.objects.db_manager("other").raw('SELECT id FROM multiple_database_book')
-        self.assertEqual(map(lambda o: o.pk, val), [dive.pk])
+        self.assertQuerysetEqual(val, [dive.pk], attrgetter("pk"))
 
         val = Book.objects.raw('SELECT id FROM multiple_database_book').using('other')
-        self.assertEqual(map(lambda o: o.pk, val), [dive.pk])
+        self.assertQuerysetEqual(val, [dive.pk], attrgetter("pk"))
 
     def test_select_related(self):
         "Database assignment is retained if an object is retrieved with select_related()"