Add initial experimental multi database support

bluetech · bluetech · commit 92c6b7ef6ee2 · 2021-05-14T11:39:28.000+03:00
diff --git a/docs/database.rst b/docs/database.rst
@@ -64,21 +64,25 @@ select using an argument to the ``django_db`` mark::
 Tests requiring multiple databases
 ----------------------------------
 
+.. caution::
+
+    This support is **experimental** and is subject to change without
+    deprecation. We are still figuring out the best way to expose this
+    functionality. If you are using this successfully or unsuccessfully,
+    `let us know <https://github.com/pytest-dev/pytest-django/issues/924>`_!
+
+``pytest-django`` has experimental support for multi-database configurations.
 Currently ``pytest-django`` does not specifically support Django's
-multi-database support.
+multi-database support, using the ``databases`` argument to the
+:py:func:`django_db <pytest.mark.django_db>` mark::
 
-You can however use normal :class:`~django.test.TestCase` instances to use its
-:ref:`django:topics-testing-advanced-multidb` support.
-In particular, if your database is configured for replication, be sure to read
-about :ref:`django:topics-testing-primaryreplica`.
+   @pytest.mark.django_db(databases=['default', 'other'])
+   def test_spam():
+       assert MyModel.objects.using('other').count() == 0
 
-If you have any ideas about the best API to support multiple databases
-directly in ``pytest-django`` please get in touch, we are interested
-in eventually supporting this but unsure about simply following
-Django's approach.
+For details see :py:attr:`django.test.TransactionTestCase.databases` and
+:py:attr:`django.test.TestCase.databases`.
 
-See `pull request 431 <https://github.com/pytest-dev/pytest-django/pull/431>`_
-for an idea/discussion to approach this.
 
 ``--reuse-db`` - reuse the testing database between test runs
 --------------------------------------------------------------
diff --git a/docs/helpers.rst b/docs/helpers.rst
@@ -24,7 +24,7 @@ Markers
 ``pytest.mark.django_db`` - request database access
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
-.. py:function:: pytest.mark.django_db([transaction=False, reset_sequences=False])
+.. py:function:: pytest.mark.django_db([transaction=False, reset_sequences=False, databases=None])
 
   This is used to mark a test function as requiring the database. It
   will ensure the database is set up correctly for the test. Each test
@@ -54,6 +54,23 @@ Markers
     effect.  Please be aware that not all databases support this feature.
     For details see :py:attr:`django.test.TransactionTestCase.reset_sequences`.
 
+
+  :type databases: Union[Iterable[str], str, None]
+  :param databases:
+    .. caution::
+
+      This argument is **experimental** and is subject to change without
+      deprecation. We are still figuring out the best way to expose this
+      functionality. If you are using this successfully or unsuccessfully,
+      `let us know <https://github.com/pytest-dev/pytest-django/issues/924>`_!
+
+    The ``databases`` argument defines which databases in a multi-database
+    configuration will be set up and may be used by the test.  Defaults to
+    only the ``default`` database.  The special value ``"__all__"`` may be use
+    to specify all configured databases.
+    For details see :py:attr:`django.test.TransactionTestCase.databases` and
+    :py:attr:`django.test.TestCase.databases`.
+
 .. note::
 
   If you want access to the Django database inside a *fixture*, this marker may
diff --git a/pytest_django/fixtures.py b/pytest_django/fixtures.py
@@ -1,5 +1,5 @@
 """All pytest-django fixtures"""
-from typing import Any, Generator, List
+from typing import Any, Generator, Iterable, List, Optional, Tuple, Union
 import os
 from contextlib import contextmanager
 from functools import partial
@@ -12,8 +12,13 @@
 
 TYPE_CHECKING = False
 if TYPE_CHECKING:
+    from typing import Literal
+
     import django
 
+    _DjangoDbDatabases = Optional[Union["Literal['__all__']", Iterable[str]]]
+    _DjangoDb = Tuple[bool, bool, _DjangoDbDatabases]
+
 
 __all__ = [
     "django_db_setup",
@@ -142,6 +147,10 @@ def _django_db_fixture_helper(
         # Do nothing, we get called with transactional=True, too.
         return
 
+    _databases = getattr(
+        request.node, "_pytest_django_databases", None,
+    )  # type: Optional[_DjangoDbDatabases]
+
     django_db_blocker.unblock()
     request.addfinalizer(django_db_blocker.restore)
 
@@ -158,6 +167,8 @@ def _django_db_fixture_helper(
     class PytestDjangoTestCase(test_case_class):  # type: ignore[misc,valid-type]
         if transactional and _reset_sequences:
             reset_sequences = True
+        if _databases is not None:
+            databases = _databases
 
     PytestDjangoTestCase.setUpClass()
     request.addfinalizer(PytestDjangoTestCase.tearDownClass)
diff --git a/pytest_django/plugin.py b/pytest_django/plugin.py
@@ -48,6 +48,8 @@
 
     import django
 
+    from .fixtures import _DjangoDb, _DjangoDbDatabases
+
 
 SETTINGS_MODULE_ENV = "DJANGO_SETTINGS_MODULE"
 CONFIGURATION_ENV = "DJANGO_CONFIGURATION"
@@ -262,12 +264,14 @@ def pytest_load_initial_conftests(
     # Register the marks
     early_config.addinivalue_line(
         "markers",
-        "django_db(transaction=False, reset_sequences=False): "
+        "django_db(transaction=False, reset_sequences=False, databases=None): "
         "Mark the test as using the Django test database.  "
         "The *transaction* argument allows you to use real transactions "
         "in the test like Django's TransactionTestCase.  "
         "The *reset_sequences* argument resets database sequences before "
-        "the test.",
+        "the test.  "
+        "The *databases* argument sets which database aliases the test "
+        "uses (by default, only 'default'). Use '__all__' for all databases.",
     )
     early_config.addinivalue_line(
         "markers",
@@ -452,7 +456,11 @@ def _django_db_marker(request) -> None:
     """
     marker = request.node.get_closest_marker("django_db")
     if marker:
-        transaction, reset_sequences = validate_django_db(marker)
+        transaction, reset_sequences, databases = validate_django_db(marker)
+
+        # TODO: Use pytest Store (item.store) once that's stable.
+        request.node._pytest_django_databases = databases
+
         if reset_sequences:
             request.getfixturevalue("django_db_reset_sequences")
         elif transaction:
@@ -727,21 +735,22 @@ def restore(self) -> None:
 _blocking_manager = _DatabaseBlocker()
 
 
-def validate_django_db(marker) -> Tuple[bool, bool]:
+def validate_django_db(marker) -> "_DjangoDb":
     """Validate the django_db marker.
 
-    It checks the signature and creates the ``transaction`` and
-    ``reset_sequences`` attributes on the marker which will have the
-    correct values.
+    It checks the signature and creates the ``transaction``,
+    ``reset_sequences`` and ``databases`` attributes on the marker
+    which will have the correct values.
 
     A sequence reset is only allowed when combined with a transaction.
     """
 
     def apifun(
         transaction: bool = False,
         reset_sequences: bool = False,
-    ) -> Tuple[bool, bool]:
-        return transaction, reset_sequences
+        databases: "_DjangoDbDatabases" = None,
+    ) -> "_DjangoDb":
+        return transaction, reset_sequences, databases
 
     return apifun(*marker.args, **marker.kwargs)
 
diff --git a/pytest_django_test/app/migrations/0001_initial.py b/pytest_django_test/app/migrations/0001_initial.py
@@ -24,5 +24,20 @@ class Migration(migrations.Migration):
                 ),
                 ("name", models.CharField(max_length=100)),
             ],
-        )
+        ),
+        migrations.CreateModel(
+            name="SecondItem",
+            fields=[
+                (
+                    "id",
+                    models.AutoField(
+                        auto_created=True,
+                        primary_key=True,
+                        serialize=False,
+                        verbose_name="ID",
+                    ),
+                ),
+                ("name", models.CharField(max_length=100)),
+            ],
+        ),
     ]
diff --git a/pytest_django_test/app/models.py b/pytest_django_test/app/models.py
@@ -1,5 +1,11 @@
 from django.db import models
 
 
+# Routed to database "main".
 class Item(models.Model):
     name = models.CharField(max_length=100)  # type: str
+
+
+# Routed to database "second".
+class SecondItem(models.Model):
+    name = models.CharField(max_length=100)  # type: str
diff --git a/pytest_django_test/db_helpers.py b/pytest_django_test/db_helpers.py
@@ -26,6 +26,9 @@
         # An explicit test db name was given, is that as the base name
         TEST_DB_NAME = "{}_inner".format(TEST_DB_NAME)
 
+SECOND_DB_NAME = DB_NAME + '_second' if DB_NAME is not None else None
+SECOND_TEST_DB_NAME = TEST_DB_NAME + '_second' if DB_NAME is not None else None
+
 
 def get_db_engine():
     return _settings["ENGINE"].split(".")[-1]
diff --git a/pytest_django_test/db_router.py b/pytest_django_test/db_router.py
@@ -0,0 +1,14 @@
+class DbRouter:
+    def db_for_read(self, model, **hints):
+        if model._meta.app_label == 'app' and model._meta.model_name == 'seconditem':
+            return 'second'
+        return None
+
+    def db_for_write(self, model, **hints):
+        if model._meta.app_label == 'app' and model._meta.model_name == 'seconditem':
+            return 'second'
+        return None
+
+    def allow_migrate(self, db, app_label, model_name=None, **hints):
+        if app_label == 'app' and model_name == 'seconditem':
+            return db == 'second'
diff --git a/pytest_django_test/settings_base.py b/pytest_django_test/settings_base.py
@@ -27,3 +27,5 @@
         "OPTIONS": {},
     }
 ]
+
+DATABASE_ROUTERS = ['pytest_django_test.db_router.DbRouter']
diff --git a/pytest_django_test/settings_mysql_innodb.py b/pytest_django_test/settings_mysql_innodb.py
@@ -6,7 +6,38 @@
 DATABASES = {
     "default": {
         "ENGINE": "django.db.backends.mysql",
-        "NAME": "pytest_django_should_never_get_accessed",
+        "NAME": "pytest_django_tests_default",
+        "USER": environ.get("TEST_DB_USER", "root"),
+        "PASSWORD": environ.get("TEST_DB_PASSWORD", ""),
+        "HOST": environ.get("TEST_DB_HOST", "localhost"),
+        "OPTIONS": {
+            "init_command": "SET default_storage_engine=InnoDB",
+            "charset": "utf8mb4",
+        },
+        "TEST": {
+            "CHARSET": "utf8mb4",
+            "COLLATION": "utf8mb4_unicode_ci",
+        },
+    },
+    "replica": {
+        "ENGINE": "django.db.backends.mysql",
+        "NAME": "pytest_django_tests_replica",
+        "USER": environ.get("TEST_DB_USER", "root"),
+        "PASSWORD": environ.get("TEST_DB_PASSWORD", ""),
+        "HOST": environ.get("TEST_DB_HOST", "localhost"),
+        "OPTIONS": {
+            "init_command": "SET default_storage_engine=InnoDB",
+            "charset": "utf8mb4",
+        },
+        "TEST": {
+            "MIRROR": "default",
+            "CHARSET": "utf8mb4",
+            "COLLATION": "utf8mb4_unicode_ci",
+        },
+    },
+    "second": {
+        "ENGINE": "django.db.backends.mysql",
+        "NAME": "pytest_django_tests_second",
         "USER": environ.get("TEST_DB_USER", "root"),
         "PASSWORD": environ.get("TEST_DB_PASSWORD", ""),
         "HOST": environ.get("TEST_DB_HOST", "localhost"),
diff --git a/pytest_django_test/settings_mysql_myisam.py b/pytest_django_test/settings_mysql_myisam.py
@@ -6,7 +6,38 @@
 DATABASES = {
     "default": {
         "ENGINE": "django.db.backends.mysql",
-        "NAME": "pytest_django_should_never_get_accessed",
+        "NAME": "pytest_django_tests_default",
+        "USER": environ.get("TEST_DB_USER", "root"),
+        "PASSWORD": environ.get("TEST_DB_PASSWORD", ""),
+        "HOST": environ.get("TEST_DB_HOST", "localhost"),
+        "OPTIONS": {
+            "init_command": "SET default_storage_engine=MyISAM",
+            "charset": "utf8mb4",
+        },
+        "TEST": {
+            "CHARSET": "utf8mb4",
+            "COLLATION": "utf8mb4_unicode_ci",
+        },
+    },
+    "replica": {
+        "ENGINE": "django.db.backends.mysql",
+        "NAME": "pytest_django_tests_replica",
+        "USER": environ.get("TEST_DB_USER", "root"),
+        "PASSWORD": environ.get("TEST_DB_PASSWORD", ""),
+        "HOST": environ.get("TEST_DB_HOST", "localhost"),
+        "OPTIONS": {
+            "init_command": "SET default_storage_engine=MyISAM",
+            "charset": "utf8mb4",
+        },
+        "TEST": {
+            "MIRROR": "default",
+            "CHARSET": "utf8mb4",
+            "COLLATION": "utf8mb4_unicode_ci",
+        },
+    },
+    "second": {
+        "ENGINE": "django.db.backends.mysql",
+        "NAME": "pytest_django_tests_second",
         "USER": environ.get("TEST_DB_USER", "root"),
         "PASSWORD": environ.get("TEST_DB_PASSWORD", ""),
         "HOST": environ.get("TEST_DB_HOST", "localhost"),
diff --git a/pytest_django_test/settings_postgres.py b/pytest_django_test/settings_postgres.py
@@ -14,7 +14,24 @@
 DATABASES = {
     "default": {
         "ENGINE": "django.db.backends.postgresql",
-        "NAME": "pytest_django_should_never_get_accessed",
+        "NAME": "pytest_django_tests_default",
+        "USER": environ.get("TEST_DB_USER", ""),
+        "PASSWORD": environ.get("TEST_DB_PASSWORD", ""),
+        "HOST": environ.get("TEST_DB_HOST", ""),
+    },
+    "replica": {
+        "ENGINE": "django.db.backends.postgresql",
+        "NAME": "pytest_django_tests_replica",
+        "USER": environ.get("TEST_DB_USER", ""),
+        "PASSWORD": environ.get("TEST_DB_PASSWORD", ""),
+        "HOST": environ.get("TEST_DB_HOST", ""),
+        "TEST": {
+            "MIRROR": "default",
+        },
+    },
+    "second": {
+        "ENGINE": "django.db.backends.postgresql",
+        "NAME": "pytest_django_tests_second",
         "USER": environ.get("TEST_DB_USER", ""),
         "PASSWORD": environ.get("TEST_DB_PASSWORD", ""),
         "HOST": environ.get("TEST_DB_HOST", ""),
diff --git a/pytest_django_test/settings_sqlite.py b/pytest_django_test/settings_sqlite.py
@@ -1,8 +1,20 @@
 from .settings_base import *  # noqa: F401 F403
 
+
 DATABASES = {
     "default": {
         "ENGINE": "django.db.backends.sqlite3",
         "NAME": "/should_not_be_accessed",
-    }
+    },
+    "replica": {
+        "ENGINE": "django.db.backends.sqlite3",
+        "NAME": "/should_not_be_accessed",
+        "TEST": {
+            "MIRROR": "default",
+        },
+    },
+    "second": {
+        "ENGINE": "django.db.backends.sqlite3",
+        "NAME": "/should_not_be_accessed",
+    },
 }
diff --git a/pytest_django_test/settings_sqlite_file.py b/pytest_django_test/settings_sqlite_file.py
diff --git a/tests/conftest.py b/tests/conftest.py
diff --git a/tests/test_database.py b/tests/test_database.py
diff --git a/tests/test_db_setup.py b/tests/test_db_setup.py

Original file line number	Diff line number	Diff line change
`@@ -27,3 +27,5 @@`
`27`	`27`	`"OPTIONS": {},`
`28`	`28`	`}`
`29`	`29`	`]`
	`30`	`+`
	`31`	`+DATABASE_ROUTERS = ['pytest_django_test.db_router.DbRouter']`