maybe a better test definition

Author: shaleh

tests/common_testing.py

@@ -3,33 +3,33 @@
 from django.test import TransactionTestCase as DjangoTransactionTestCase
 
 
-class OAuth2ProviderTestCase(DjangoTestCase):
-    """Place holder to allow overriding behaviors."""
+# When there are multiple databases defined, Django tests will not work unless they are
+# told which database(s) to work with. The multiple database scenario setup for these
+# tests purposefully defines 'default' as an empty database in order to catch any
+# assumptions in this package about database names and in particular to ensure there is
+# no assumption that 'default' is a valid database.
+# For any test that would usually use Django's TestCase or TransactionTestCase using
+# the classes defined here is all that is required.
+# Any test that uses pytest's django_db need to include a databases parameter using
+# test_database_names defined below.
+# In test code, anywhere the database is referenced the Django router needs to be used
+# exactly like the package's code.
+# For instance:
+#     token_database = router.db_for_write(AccessToken)
+#     with self.assertNumQueries(1, using=token_database):
+# Without the 'using' option, this test fails in the multiple database scenario because
+# 'default' is used.
+
+test_database_names = ["alpha", "beta"] if len(settings.DATABASES) > 1 else ["default"]
+
 
+class OAuth2ProviderBase:
+    databases = test_database_names
 
-class OAuth2ProviderTransactionTestCase(DjangoTransactionTestCase):
+
+class OAuth2ProviderTestCase(OAuth2ProviderBase, DjangoTestCase):
     """Place holder to allow overriding behaviors."""
 
 
-if len(settings.DATABASES) > 1:
-    # There are multiple databases defined. When this happens Django tests will not
-    # work unless they are told which database(s) to work with. The multiple
-    # database scenario setup for these tests purposefully defines 'default' as an
-    # empty database in order to catch any assumptions in this package about database
-    # names and in particular to ensure there is no assumption that 'default' is a
-    # valid database.
-    # For any test that would usually use Django's TestCase or TransactionTestCase
-    # using the classes defined here is all that is required.
-    # Any test that uses pytest's django_db need to base in a databases parameter
-    # using this definition of test_database_names.
-    # In test code, anywhere the default database is used the variable
-    # database_for_oauth2_provider must be used in its place. For instance,
-    #     with self.assertNumQueries(1, using=database_for_oauth2_provider):
-    # without the using option this fails because default is used.
-    test_database_names = {name for name in settings.DATABASES if name != "default"}
-    database_for_oauth2_provider = "alpha"
-    OAuth2ProviderTestCase.databases = test_database_names
-    OAuth2ProviderTransactionTestCase.databases = test_database_names
-else:
-    test_database_names = {"default"}
-    database_for_oauth2_provider = "default"
+class OAuth2ProviderTransactionTestCase(OAuth2ProviderBase, DjangoTransactionTestCase):
+    """Place holder to allow overriding behaviors."""

tests/test_introspection_view.py

@@ -3,14 +3,14 @@
 
 import pytest
 from django.contrib.auth import get_user_model
+from django.db import router
 from django.urls import reverse
 from django.utils import timezone
 
 from oauth2_provider.models import get_access_token_model, get_application_model
 
 from . import presets
 from .common_testing import OAuth2ProviderTestCase as TestCase
-from .common_testing import database_for_oauth2_provider
 from .utils import get_basic_auth_header
 
 
@@ -344,5 +344,6 @@ def test_view_post_invalid_client_creds_plaintext(self):
         self.assertEqual(response.status_code, 403)
 
     def test_select_related_in_view_for_less_db_queries(self):
-        with self.assertNumQueries(1, using=database_for_oauth2_provider):
+        token_database = router.db_for_write(AccessToken)
+        with self.assertNumQueries(1, using=token_database):
             self.client.post(reverse("oauth2_provider:introspect"))