diff --git a/django_prometheus/migrations.py b/django_prometheus/migrations.py index c5ca6188..e264db36 100644 --- a/django_prometheus/migrations.py +++ b/django_prometheus/migrations.py @@ -38,13 +38,14 @@ def ExportMigrations(): # ready. from django.db.migrations.executor import MigrationExecutor - if "default" in connections and (isinstance(connections["default"], DatabaseWrapper)): - # This is the case where DATABASES = {} in the configuration, - # i.e. the user is not using any databases. Django "helpfully" - # adds a dummy database and then throws when you try to - # actually use it. So we don't do anything, because trying to - # export stats would crash the app on startup. - return for alias in connections.databases: + if alias == "default" and (isinstance(connections[alias], DatabaseWrapper)): + # This is the case where DATABASES = {} in the configuration, + # i.e. the user is not using any databases. Django "helpfully" + # adds a dummy database and then throws when you try to + # actually use it. So we don't do anything, because trying to + # export stats would crash the app on startup. + continue + executor = MigrationExecutor(connections[alias]) ExportMigrationsForDatabase(alias, executor) diff --git a/django_prometheus/tests/end2end/testapp/test_migrations.py b/django_prometheus/tests/end2end/testapp/test_migrations.py index 26e46561..61f34d27 100644 --- a/django_prometheus/tests/end2end/testapp/test_migrations.py +++ b/django_prometheus/tests/end2end/testapp/test_migrations.py @@ -1,8 +1,8 @@ -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch import pytest -from django_prometheus.migrations import ExportMigrationsForDatabase +from django_prometheus.migrations import ExportMigrations, ExportMigrationsForDatabase from django_prometheus.testutils import assert_metric_equal @@ -35,3 +35,76 @@ def test_counters(self): assert_metric_equal(0, M("unapplied_total"), connection="fakedb1") assert_metric_equal(2, M("applied_total"), connection="fakedb2") assert_metric_equal(1, M("unapplied_total"), connection="fakedb2") + + +class TestExportMigrations: + """Test ExportMigrations function with mocked connections.""" + + @patch("django_prometheus.migrations.connections") + @patch("django.db.migrations.executor.MigrationExecutor") + def test_export_migrations_with_real_database(self, mock_executor_class, mock_connections): + """Test ExportMigrations with a non-empty database connection.""" + + # Create a mock connection that is NOT a DatabaseWrapper (dummy) + mock_connections.databases = {"default": {}} + mock_connection = MagicMock() + mock_connection.__class__.__name__ = "RealDatabaseWrapper" + mock_connections.__getitem__.return_value = mock_connection + + executor = MagicMock() + executor.migration_plan.return_value = ["migration1", "migration2"] + executor.loader.applied_migrations = {"m1", "m2", "m3"} + mock_executor_class.return_value = executor + + ExportMigrations() + + # Verify MigrationExecutor was called with the connection + mock_executor_class.assert_called_once_with(mock_connection) + assert executor.migration_plan.call_count == 1 + + @patch("django_prometheus.migrations.connections") + @patch("django.db.migrations.executor.MigrationExecutor") + def test_export_migrations_skips_dummy_database(self, mock_executor_class, mock_connections): + """Test ExportMigrations skips dummy database (DATABASES = {} case).""" + from django.db.backends.dummy.base import DatabaseWrapper + + # Create a mock connection that IS a DatabaseWrapper (dummy) + mock_connections.databases = {"default": {}} + mock_dummy_connection = MagicMock(spec=DatabaseWrapper) + mock_connections.__getitem__.return_value = mock_dummy_connection + + ExportMigrations() + + # Verify MigrationExecutor was NOT called (because we skip dummy databases) + mock_executor_class.assert_not_called() + + @patch("django_prometheus.migrations.connections") + @patch("django.db.migrations.executor.MigrationExecutor") + def test_export_migrations_with_multiple_databases(self, mock_executor_class, mock_connections): + """Test ExportMigrations with multiple databases including a dummy one.""" + from django.db.backends.dummy.base import DatabaseWrapper + + mock_connections.databases = {"default": {}, "secondary": {}} + + mock_dummy_connection = MagicMock(spec=DatabaseWrapper) + mock_real_connection = MagicMock() + mock_real_connection.__class__.__name__ = "RealDatabaseWrapper" + + # Setup the connections dict to return different connections for different aliases + def get_connection(alias): + if alias == "default": + return mock_dummy_connection + else: + return mock_real_connection + + mock_connections.__getitem__.side_effect = get_connection + + executor = MagicMock() + executor.migration_plan.return_value = [] + executor.loader.applied_migrations = set() + mock_executor_class.return_value = executor + + ExportMigrations() + + # Verify MigrationExecutor was called only once (for the non-dummy database) + mock_executor_class.assert_called_once_with(mock_real_connection)