Run django db migrations for tests (#729)

Author: michelletran-sentry

conftest.py

@@ -9,9 +9,8 @@
 from shared.storage.memory import MemoryStorageService
 from shared.torngit import Github as GithubHandler
 from sqlalchemy import event
-from sqlalchemy.exc import OperationalError
 from sqlalchemy.orm import Session
-from sqlalchemy_utils import create_database, database_exists
+from sqlalchemy_utils import database_exists
 
 from celery_config import initialize_logging
 from database.base import Base
@@ -40,7 +39,7 @@ def pytest_itemcollected(item):
 
 
 @pytest.fixture(scope="session")
-def engine(request, sqlalchemy_connect_url, app_config):
+def engine(request, sqlalchemy_db, sqlalchemy_connect_url, app_config):
     """Engine configuration.
     See http://docs.sqlalchemy.org/en/latest/core/engines.html
     for more details.
@@ -67,30 +66,87 @@ def engine(request, sqlalchemy_connect_url, app_config):
         engine.url.database = "{}_{}".format(engine.url.database, xdist_suffix)
         engine = create_engine(engine.url)  # override engine
 
-    def fin():
-        print("Disposing engine")
-        engine.dispose()
+    # Check that the DB exist and migrate the unmigrated SQLALchemy models as a stop-gap
+    database_url = sqlalchemy_connect_url
+    if not database_exists(database_url):
+        raise RuntimeError(f"SQLAlchemy cannot connect to DB at {database_url}")
+
+    Base.metadata.tables["profiling_profilingcommit"].create(bind=engine)
+    Base.metadata.tables["profiling_profilingupload"].create(bind=engine)
+    Base.metadata.tables["timeseries_measurement"].create(bind=engine)
+    Base.metadata.tables["timeseries_dataset"].create(bind=engine)
+
+    Base.metadata.tables["compare_commitcomparison"].create(bind=engine)
+    Base.metadata.tables["compare_flagcomparison"].create(bind=engine)
+    Base.metadata.tables["compare_componentcomparison"].create(bind=engine)
+
+    Base.metadata.tables["labelanalysis_labelanalysisrequest"].create(bind=engine)
+    Base.metadata.tables["labelanalysis_labelanalysisprocessingerror"].create(
+        bind=engine
+    )
+
+    Base.metadata.tables["staticanalysis_staticanalysissuite"].create(bind=engine)
+    Base.metadata.tables["staticanalysis_staticanalysissinglefilesnapshot"].create(
+        bind=engine
+    )
+    Base.metadata.tables["staticanalysis_staticanalysissuitefilepath"].create(
+        bind=engine
+    )
+
+    yield engine
 
-    request.addfinalizer(fin)
-    return engine
+    print("Disposing engine")
+    engine.dispose()
 
 
 @pytest.fixture(scope="session")
-def db(engine, sqlalchemy_connect_url):
-    database_url = sqlalchemy_connect_url
+def sqlalchemy_db(request: pytest.FixtureRequest, django_db_blocker, django_db_setup):
+    # Bootstrap the DB by running the Django bootstrap version.
+    from django.conf import settings
+    from django.db import connections
+    from django.test.utils import setup_databases, teardown_databases
+
+    with django_db_blocker.unblock():
+        # Temporarily reset the database to the SQLAlchemy DBs to run the migrations.
+        original_test_name = settings.DATABASES["default"]["TEST"]["NAME"]
+        settings.DATABASES["default"]["TEST"]["NAME"] = "test_postgres_sqlalchemy"
+        db_cfg = setup_databases(
+            verbosity=request.config.option.verbose,
+            interactive=False,
+        )
+        settings.DATABASES["default"]["TEST"]["NAME"] = original_test_name
+
+        # Hack to get the default connection for the test database to _actually_ be the
+        # Django database that the django_db should actually use. It was set to the SQLAlchemy database,
+        # but this makes sure that the default Django DB connection goes to the Django database.
+        # Since the database was already created and migrated in the django_db_setup fixture,
+        # we set keepdb=True to avoid recreating the database and rerunning the migrations.
+        connections.configure_settings(settings.DATABASES)
+        connections["default"].creation.create_test_db(
+            verbosity=request.config.option.verbose,
+            autoclobber=True,
+            keepdb=True,
+        )
+
+    yield
+
+    # Cleanup with Django version as well
     try:
-        if not database_exists(database_url):
-            create_database(database_url)
-    except OperationalError:
-        pytest.skip("No available db")
-    connection = engine.connect()
-    connection.execute("DROP SCHEMA IF EXISTS public CASCADE;")
-    connection.execute("CREATE SCHEMA public;")
-    Base.metadata.create_all(engine)
+        with django_db_blocker.unblock():
+            settings.DATABASES["default"]["TEST"]["NAME"] = "test_postgres_sqlalchemy"
+            teardown_databases(db_cfg, verbosity=request.config.option.verbose)
+            settings.DATABASES["default"]["TEST"]["NAME"] = original_test_name
+    except Exception as exc:  # noqa: BLE001
+        request.node.warn(
+            pytest.PytestWarning(
+                f"Error when trying to teardown test databases: {exc!r}"
+            )
+        )
 
 
 @pytest.fixture
-def dbsession(db, engine):
+def dbsession(sqlalchemy_db, engine):
+    """Sets up the SQLAlchemy dbsession."""
     connection = engine.connect()
 
     connection_transaction = connection.begin()
@@ -239,7 +295,7 @@ def side_effect(*args, **kwargs):
 @pytest.fixture
 def with_sql_functions(dbsession):
     dbsession.execute(
-        """CREATE FUNCTION array_append_unique(anyarray, anyelement) RETURNS anyarray
+        """CREATE OR REPLACE FUNCTION array_append_unique(anyarray, anyelement) RETURNS anyarray
                 LANGUAGE sql IMMUTABLE
                 AS $_$
             select case when $2 is null

database/tests/factories/core.py

@@ -1,4 +1,4 @@
-from datetime import datetime
+from datetime import datetime, timezone
 from hashlib import sha1
 from uuid import uuid4
 
@@ -71,9 +71,11 @@ class Meta:
 
     owner = factory.SubFactory(OwnerFactory)
     bot = None
-    updatestamp = factory.LazyAttribute(lambda o: datetime.now())
+    updatestamp = factory.LazyAttribute(lambda o: datetime.now(tz=timezone.utc))
     languages = []
-    languages_last_updated = factory.LazyAttribute(lambda o: datetime.now())
+    languages_last_updated = factory.LazyAttribute(
+        lambda o: datetime.now(tz=timezone.utc)
+    )
     bundle_analysis_enabled = False
 
 
@@ -113,7 +115,7 @@ class Meta:
     )
     ci_passed = True
     pullid = None
-    timestamp = datetime(2019, 2, 1, 17, 59, 47)
+    timestamp = datetime(2019, 2, 1, 17, 59, 47, tzinfo=timezone.utc)
     author = factory.SubFactory(OwnerFactory)
     repository = factory.SubFactory(RepositoryFactory)
     totals = factory.LazyFunction(

database/tests/unit/test_model_utils.py

@@ -56,14 +56,14 @@ def test_subclass_validation(self, mocker):
             self.ClassWithArchiveFieldMissingMethods, ArchiveFieldInterface
         )
 
-    def test_archive_getter_db_field_set(self, db):
+    def test_archive_getter_db_field_set(self, sqlalchemy_db):
         commit = CommitFactory()
         test_class = self.ClassWithArchiveField(commit, "db_value", "gcs_path")
         assert test_class._archive_field == "db_value"
         assert test_class._archive_field_storage_path == "gcs_path"
         assert test_class.archive_field == "db_value"
 
-    def test_archive_getter_archive_field_set(self, db, mocker):
+    def test_archive_getter_archive_field_set(self, sqlalchemy_db, mocker):
         some_json = {"some": "data"}
         mock_read_file = mocker.MagicMock(return_value=json.dumps(some_json))
         mock_archive_service = mocker.patch("database.utils.ArchiveService")
@@ -81,7 +81,7 @@ def test_archive_getter_archive_field_set(self, db, mocker):
         assert test_class.archive_field == some_json
         assert mock_read_file.call_count == 1
 
-    def test_archive_getter_file_not_in_storage(self, db, mocker):
+    def test_archive_getter_file_not_in_storage(self, sqlalchemy_db, mocker):
         mocker.patch(
             "database.utils.ArchiveField.read_timeout",
             new_callable=PropertyMock,
@@ -99,7 +99,7 @@ def test_archive_getter_file_not_in_storage(self, db, mocker):
         mock_read_file.assert_called_with("gcs_path")
         mock_archive_service.assert_called_with(repository=commit.repository)
 
-    def test_archive_setter_db_field(self, db, mocker):
+    def test_archive_setter_db_field(self, sqlalchemy_db, mocker):
         commit = CommitFactory()
         test_class = self.ClassWithArchiveField(commit, "db_value", "gcs_path", False)
         assert test_class._archive_field == "db_value"
@@ -111,7 +111,7 @@ def test_archive_setter_db_field(self, db, mocker):
         assert test_class._archive_field == "batata frita"
         assert test_class.archive_field == "batata frita"
 
-    def test_archive_setter_archive_field(self, db, mocker):
+    def test_archive_setter_archive_field(self, sqlalchemy_db, mocker):
         commit = CommitFactory()
         test_class = self.ClassWithArchiveField(commit, "db_value", None, True)
         some_json = {"some": "data"}

django_scaffold/tests_settings.py

@@ -22,3 +22,5 @@
         },
     }
 ]
+
+DATABASES["default"]["TEST"] = {"NAME": "test_postgres_django"}

helpers/tests/unit/test_github_installation.py

@@ -18,7 +18,7 @@ def test_get_installation_name_for_owner_for_task(dbsession: Session):
         installation_name="my_installation",
         task_name=task_name,
     )
-    dbsession.add_all([owner, installation_task_config])
+    dbsession.add_all([owner, other_owner, installation_task_config])
     dbsession.flush()
     assert (
         get_installation_name_for_owner_for_task(task_name, owner) == "my_installation"

one_off_scripts/tests/test_backfill_daily_test_rollups.py

@@ -1,85 +1,17 @@
 import datetime as dt
 
 import pytest
-import time_machine
-from shared.django_apps.core.tests.factories import RepositoryFactory
 from shared.django_apps.reports.models import DailyTestRollup
-from shared.django_apps.reports.tests.factories import (
-    DailyTestRollupFactory,
-    FlakeFactory,
-    TestFactory,
-    TestInstanceFactory,
-)
 
 from one_off_scripts.backfill_daily_test_rollups import backfill_test_rollups
-
-
-@pytest.fixture
-def setup_tests(transactional_db):
-    repo_1 = RepositoryFactory(test_analytics_enabled=True)
-
-    test_1 = TestFactory(repository=repo_1)
-
-    _ = FlakeFactory(
-        test=test_1,
-        repository=repo_1,
-        start_date=dt.datetime.fromisoformat("1970-01-02T00:00:00Z"),
-        end_date=dt.datetime.fromisoformat("1970-01-04T00:00:00Z"),
-    )
-
-    _ = FlakeFactory(
-        test=test_1,
-        repository=repo_1,
-        start_date=dt.datetime.fromisoformat("1970-01-04T12:00:00Z"),
-        end_date=None,
-    )
-
-    traveller = time_machine.travel("1970-01-01T00:00:00Z", tick=False)
-    traveller.start()
-    ti = TestInstanceFactory(test=test_1, duration_seconds=10.0)
-    traveller.stop()
-
-    traveller = time_machine.travel("1970-01-03T00:00:00Z", tick=False)
-    traveller.start()
-    ti = TestInstanceFactory(test=test_1, duration_seconds=10.0)
-    traveller.stop()
-
-    traveller = time_machine.travel("1970-01-05T00:00:00Z", tick=False)
-    traveller.start()
-    ti = TestInstanceFactory(
-        test=test_1,
-        duration_seconds=10000.0,
-    )
-    traveller.stop()
-
-    _ = DailyTestRollupFactory(
-        test=test_1,
-        date=dt.date.fromisoformat("1970-01-03"),
-        fail_count=10,
-        pass_count=5,
-        last_duration_seconds=10.0,
-        avg_duration_seconds=1.0,
-        latest_run=dt.datetime.fromisoformat("1970-01-01T00:00:00Z"),
-    )
-
-    _ = DailyTestRollupFactory(
-        test=test_1,
-        date=dt.date.fromisoformat("1970-01-05"),
-        fail_count=10,
-        pass_count=5,
-        last_duration_seconds=10.0,
-        avg_duration_seconds=1.0,
-        latest_run=dt.datetime.fromisoformat("1970-01-01T00:00:00Z"),
-    )
-
-    _ = RepositoryFactory(test_analytics_enabled=False)
-
-    return repo_1
+from one_off_scripts.tests.utils import setup_one_off_tests
 
 
 @pytest.mark.django_db(transaction=True)
-def test_backfill_test_rollups(setup_tests):
-    rollups = DailyTestRollup.objects.all()
+def test_backfill_test_rollups():
+    setup_one_off_tests()
+
+    rollups = DailyTestRollup.objects.all().order_by("date")
     assert [
         {
             "date": r.date,
@@ -123,7 +55,7 @@ def test_backfill_test_rollups(setup_tests):
         end_date="1970-01-04",
     )
 
-    rollups = DailyTestRollup.objects.all()
+    rollups = DailyTestRollup.objects.all().order_by("date")
 
     assert [
         {
@@ -138,16 +70,6 @@ def test_backfill_test_rollups(setup_tests):
         }
         for r in rollups
     ] == [
-        {
-            "avg_duration_seconds": 1.0,
-            "date": dt.date(1970, 1, 5),
-            "fail_count": 10,
-            "flaky_fail_count": 0,
-            "last_duration_seconds": 10.0,
-            "latest_run": dt.datetime.fromisoformat("1970-01-01T00:00:00Z"),
-            "pass_count": 5,
-            "skip_count": 0,
-        },
         {
             "avg_duration_seconds": 10.0,
             "date": dt.date(1970, 1, 3),
@@ -158,16 +80,26 @@ def test_backfill_test_rollups(setup_tests):
             "pass_count": 0,
             "skip_count": 0,
         },
+        {
+            "avg_duration_seconds": 1.0,
+            "date": dt.date(1970, 1, 5),
+            "fail_count": 10,
+            "flaky_fail_count": 0,
+            "last_duration_seconds": 10.0,
+            "latest_run": dt.datetime.fromisoformat("1970-01-01T00:00:00Z"),
+            "pass_count": 5,
+            "skip_count": 0,
+        },
     ]
 
-    assert len(rollups[1].commits_where_fail) == 1
+    assert len(rollups[0].commits_where_fail) == 1
 
     backfill_test_rollups(
         start_date="1970-01-02",
         end_date="1970-01-05",
     )
 
-    rollups = DailyTestRollup.objects.all()
+    rollups = DailyTestRollup.objects.all().order_by("date")
 
     assert [
         {
@@ -203,6 +135,5 @@ def test_backfill_test_rollups(setup_tests):
             "skip_count": 0,
         },
     ]
-
     for r in rollups:
         assert len(r.commits_where_fail) == 1