from django.db import connection, migrations, models
from django.db.migrations.state import ProjectState
from django.test import override_settings
from .test_base import OperationTestBase
class AgnosticRouter:
"""
A router that doesn't have an opinion regarding migrating."""def allow_migrate(self, db, app_label, **hints):
return None
class MigrateNothingRouter:
"""
A router that doesn't allow migrating."""def allow_migrate(self, db, app_label, **hints):
return False
class MigrateEverythingRouter:
"""
A router that always allows migrating."""def allow_migrate(self, db, app_label, **hints):
return True
class MigrateWhenFooRouter:
"""
A router that allows migrating depending on a hint."""def allow_migrate(self, db, app_label, **hints):
return hints.get("foo", False)
class MultiDBOperationTests(OperationTestBase):
databases = {"default", "other"}
def _test_create_model(self, app_label, should_run):
"""
CreateModel honors multi-db settings."""operation = migrations.CreateModel(
"Pony",
[("id", models.AutoField(primary_key=True))],
)# Test the state alteration
project_state = ProjectState()
new_state = project_state.clone()
operation.state_forwards(app_label, new_state)
# Test the database alteration
self.assertTableNotExists("%s_pony" % app_label)
with connection.schema_editor() as editor:
operation.database_forwards(app_label, editor, project_state, new_state)
if should_run:
self.assertTableExists("%s_pony" % app_label)
else:
self.assertTableNotExists("%s_pony" % app_label)
# And test reversal
with connection.schema_editor() as editor:
operation.database_backwards(app_label, editor, new_state, project_state)
self.assertTableNotExists("%s_pony" % app_label)
@override_settings(DATABASE_ROUTERS=[AgnosticRouter()])
def test_create_model(self):
"""
Test when router doesn't have an opinion (i.e. CreateModel should run)."""self._test_create_model("test_mltdb_crmo", should_run=True)
@override_settings(DATABASE_ROUTERS=[MigrateNothingRouter()])
def test_create_model2(self):
"""
Test when router returns False (i.e. CreateModel shouldn't run)."""self._test_create_model("test_mltdb_crmo2", should_run=False)
@override_settings(DATABASE_ROUTERS=[MigrateEverythingRouter()])
def test_create_model3(self):
"""
Test when router returns True (i.e. CreateModel should run)."""self._test_create_model("test_mltdb_crmo3", should_run=True)
def test_create_model4(self):
"""
Test multiple routers."""with override_settings(DATABASE_ROUTERS=[AgnosticRouter(), AgnosticRouter()]):
self._test_create_model("test_mltdb_crmo4", should_run=True)
with override_settings(
DATABASE_ROUTERS=[MigrateNothingRouter(), MigrateEverythingRouter()]
):self._test_create_model("test_mltdb_crmo4", should_run=False)
with override_settings(
DATABASE_ROUTERS=[MigrateEverythingRouter(), MigrateNothingRouter()]
):self._test_create_model("test_mltdb_crmo4", should_run=True)
def _test_run_sql(self, app_label, should_run, hints=None):
with override_settings(DATABASE_ROUTERS=[MigrateEverythingRouter()]):
project_state = self.set_up_test_model(app_label)
sql = """INSERT INTO {0}_pony (pink, weight) VALUES (1, 3.55);
INSERT INTO {0}_pony (pink, weight) VALUES (3, 5.0);
""".format(
app_label)operation = migrations.RunSQL(sql, hints=hints or {})
# Test the state alteration does nothing
new_state = project_state.clone()
operation.state_forwards(app_label, new_state)
self.assertEqual(new_state, project_state)
# Test the database alteration
self.assertEqual(
project_state.apps.get_model(app_label, "Pony").objects.count(), 0
)with connection.schema_editor() as editor:
operation.database_forwards(app_label, editor, project_state, new_state)
Pony = project_state.apps.get_model(app_label, "Pony")
if should_run:
self.assertEqual(Pony.objects.count(), 2)
else:
self.assertEqual(Pony.objects.count(), 0)
@override_settings(DATABASE_ROUTERS=[MigrateNothingRouter()])
def test_run_sql_migrate_nothing_router(self):
self._test_run_sql("test_mltdb_runsql", should_run=False)
@override_settings(DATABASE_ROUTERS=[MigrateWhenFooRouter()])
def test_run_sql_migrate_foo_router_without_hints(self):
self._test_run_sql("test_mltdb_runsql2", should_run=False)
@override_settings(DATABASE_ROUTERS=[MigrateWhenFooRouter()])
def test_run_sql_migrate_foo_router_with_hints(self):
self._test_run_sql("test_mltdb_runsql3", should_run=True, hints={"foo": True})
def _test_run_python(self, app_label, should_run, hints=None):
with override_settings(DATABASE_ROUTERS=[MigrateEverythingRouter()]):
project_state = self.set_up_test_model(app_label)
# Create the operation
def inner_method(models, schema_editor):
Pony = models.get_model(app_label, "Pony")
Pony.objects.create(pink=1, weight=3.55)
Pony.objects.create(weight=5)
operation = migrations.RunPython(inner_method, hints=hints or {})
# Test the state alteration does nothing
new_state = project_state.clone()
operation.state_forwards(app_label, new_state)
self.assertEqual(new_state, project_state)
# Test the database alteration
self.assertEqual(
project_state.apps.get_model(app_label, "Pony").objects.count(), 0
)with connection.schema_editor() as editor:
operation.database_forwards(app_label, editor, project_state, new_state)
Pony = project_state.apps.get_model(app_label, "Pony")
if should_run:
self.assertEqual(Pony.objects.count(), 2)
else:
self.assertEqual(Pony.objects.count(), 0)
@override_settings(DATABASE_ROUTERS=[MigrateNothingRouter()])
def test_run_python_migrate_nothing_router(self):
self._test_run_python("test_mltdb_runpython", should_run=False)
@override_settings(DATABASE_ROUTERS=[MigrateWhenFooRouter()])
def test_run_python_migrate_foo_router_without_hints(self):
self._test_run_python("test_mltdb_runpython2", should_run=False)
@override_settings(DATABASE_ROUTERS=[MigrateWhenFooRouter()])
def test_run_python_migrate_foo_router_with_hints(self):
self._test_run_python(
"test_mltdb_runpython3", should_run=True, hints={"foo": True}
)