feat(dp): migration tests framework + poc test

dp/cloud/python/magma/configuration_controller/tests/unit/test_request_consumer.py

@@ -1,6 +1,7 @@
 from magma.configuration_controller.request_consumer.request_db_consumer import (
     RequestDBConsumer,
 )
+from magma.db_service.config import TestConfig
 from magma.db_service.models import (
     DBCbsd,
     DBCbsdState,
@@ -56,7 +57,7 @@ def test_different_processes_dont_pick_up_each_others_requests(self, max_batch_s
         that have no lock on them.
         """
         # Given
-        config = self.get_config()
+        config = TestConfig()
         config.REQUEST_PROCESSING_LIMIT = max_batch_size
         session1 = Session(bind=self.engine)
         session2 = Session(bind=self.engine)

dp/cloud/python/magma/db_service/migrations/alembic.ini

@@ -39,8 +39,6 @@ prepend_sys_path = .
 # are written from script.py.mako
 # output_encoding = utf-8
 
-sqlalchemy.url = os.environ.get('SQLALCHEMY_DB_URI')
-
 
 [post_write_hooks]
 # post_write_hooks defines scripts or Python functions that are run

dp/cloud/python/magma/db_service/migrations/env.py

@@ -33,11 +33,15 @@
 # other values from the config, defined by the needs of env.py,
 # can be acquired:
 # ... etc.
-db_service_cfg = get_config()
-db_url = db_service_cfg.SQLALCHEMY_DB_URI
 
-if db_url:
-    config.set_section_option("alembic", "sqlalchemy.url", db_url)
+db_url = config.get_section("alembic").get("sqlalchemy.url")
+
+if not db_url:
+    db_service_cfg = get_config()
+    db_url = db_service_cfg.SQLALCHEMY_DB_URI
+
+    if db_url:
+        config.set_section_option("alembic", "sqlalchemy.url", db_url)
 
 
 def run_migrations_offline():

dp/cloud/python/magma/db_service/tests/alembic_testcase.py

@@ -0,0 +1,75 @@
+import os
+from typing import List, Optional
+
+import alembic.config
+import sqlalchemy
+import testing.postgresql
+from alembic.command import downgrade, stamp, upgrade
+from alembic.script import ScriptDirectory
+from magma.db_service.tests.db_testcase import DBTestCaseBlueprint
+
+HERE = os.path.abspath(os.path.dirname(__file__))
+PROJECT_ROOT = os.path.join(HERE, os.pardir)
+
+
+class AlembicTestCase(DBTestCaseBlueprint):
+    postgresql: testing.postgresql.Postgresql
+    up_revision: Optional[str] = None
+    down_revision: Optional[str] = None
+    tables: Optional[List[str]] = None
+
+    def tearDown(self) -> None:
+        super().tearDown()
+        self.close_session()
+        self.reflect_tables()
+        self.drop_all()
+        stamp(self.alembic_config, ())
+
+    @classmethod
+    def setUpClass(cls) -> None:
+        cls.postgresql = testing.postgresql.PostgresqlFactory(cache_initialized_db=True)()
+
+    @classmethod
+    def tearDownClass(cls) -> None:
+        cls.postgresql.stop()
+
+    def setUp(self) -> None:
+        self.metadata = sqlalchemy.MetaData()
+        self.set_up_db_test_case(SQLALCHEMY_DB_URI=self.postgresql.url())
+        self.up_revision = self.up_revision or "head"
+        self.down_revision = self.down_revision or "base"
+        self.alembic_config = alembic.config.Config(os.path.join(PROJECT_ROOT, 'migrations/alembic.ini'))
+        self.alembic_config.set_section_option('alembic', 'sqlalchemy.url', self.postgresql.url())
+        self.alembic_config.set_main_option('script_location', os.path.join(PROJECT_ROOT, 'migrations'))
+        self.script = ScriptDirectory.from_config(self.alembic_config)
+        # Making sure there are no tables in metadata
+        self.assertListEqual(self.metadata.sorted_tables, [])
+
+    def get_table(self, table_name):
+        return sqlalchemy.Table(table_name, sqlalchemy.MetaData(), autoload_with=self.engine)
+
+    def upgrade(self, revision=None):
+        revision = revision or self.up_revision
+        upgrade(self.alembic_config, revision=revision)
+
+    def downgrade(self, revision=None):
+        revision = revision or self.down_revision
+        downgrade(self.alembic_config, revision=revision)
+
+    def reflect_tables(self):
+        self.metadata.reflect()
+
+    def close_session(self):
+        self.session.rollback()
+        self.session.close()
+
+    def drop_all(self):
+        self.metadata.drop_all()
+
+    def then_tables_are(self, tables=None):
+        if tables is None:
+            tables = []
+        self.reflect_tables()
+        metadata_tables = dict(self.metadata.tables)
+        del metadata_tables["alembic_version"]
+        self.assertListEqual(sorted(tables), sorted(metadata_tables.keys()))

dp/cloud/python/magma/db_service/tests/db_testcase.py

@@ -1,33 +1,54 @@
 import unittest
+from typing import Dict, Optional
 
+import sqlalchemy.engine
 from magma.db_service.config import TestConfig
 from magma.db_service.models import Base
 from magma.db_service.session_manager import Session
-from sqlalchemy import create_engine
+from sqlalchemy import MetaData, create_engine
 
 
-class DBTestCase(unittest.TestCase):
+class DBTestCaseBlueprint(unittest.TestCase):
+    metadata: MetaData
+    engine: sqlalchemy.engine.Engine
+    session: Session
 
-    def get_config(self):
-        return TestConfig()
+    def drop_all(self):
+        self.metadata.drop_all()
 
-    def setUp(self):
-        config = self.get_config()
-        self.engine = create_engine(
-            url=config.SQLALCHEMY_DB_URI,
-            encoding=config.SQLALCHEMY_DB_ENCODING,
+    def set_up_db_test_case(self, **kwargs: Optional[Dict]):
+        self.engine = self.get_test_db_engine(**kwargs)
+        self.session = Session(bind=self.engine)
+        self.bind_engine()
+
+    @staticmethod
+    def get_test_db_engine(**kwargs) -> sqlalchemy.engine.Engine:
+        config = TestConfig()
+        return create_engine(
+            url=kwargs.get("SQLALCHEMY_DB_URI") or config.SQLALCHEMY_DB_URI,
+            encoding=kwargs.get("SQLALCHEMY_DB_ENCODING") or config.SQLALCHEMY_DB_ENCODING,
             echo=False,
-            future=config.SQLALCHEMY_FUTURE,
+            future=kwargs.get("SQLALCHEMY_FUTURE") or config.SQLALCHEMY_FUTURE,
         )
-        Base.metadata.bind = self.engine
-        Base.metadata.create_all()
-        self.session = Session()
 
-    def tearDown(self):
+    def bind_engine(self):
+        self.metadata.bind = self.engine
+
+    def close_session(self):
         self.session.rollback()
         self.session.close()
+
+
+class BaseDBTestCase(DBTestCaseBlueprint):
+
+    def setUp(self):
+        self.metadata = Base.metadata
+        self.set_up_db_test_case()
+        self.create_all()
+
+    def tearDown(self):
+        self.close_session()
         self.drop_all()
 
-    @staticmethod
-    def drop_all():
-        Base.metadata.drop_all()
+    def create_all(self):
+        self.metadata.create_all()

dp/cloud/python/magma/db_service/tests/local_db_test_case.py

@@ -1,21 +1,20 @@
 import testing.postgresql
-from magma.db_service.tests.db_testcase import DBTestCase
+from magma.db_service.models import Base
+from magma.db_service.tests.db_testcase import BaseDBTestCase
 
-Postgresql = testing.postgresql.PostgresqlFactory(cache_initialized_db=True)
 
-
-class LocalDBTestCase(DBTestCase):
+class LocalDBTestCase(BaseDBTestCase):
     postgresql: testing.postgresql.Postgresql
 
     @classmethod
     def setUpClass(cls) -> None:
-        cls.postgresql = Postgresql()
+        cls.postgresql = testing.postgresql.PostgresqlFactory(cache_initialized_db=True)()
 
     @classmethod
     def tearDownClass(cls) -> None:
         cls.postgresql.stop()
 
-    def get_config(self):
-        config = super().get_config()
-        config.SQLALCHEMY_DB_URI = self.postgresql.url()
-        return config
+    def setUp(self):
+        self.metadata = Base.metadata
+        self.set_up_db_test_case(SQLALCHEMY_DB_URI=self.postgresql.url())
+        self.create_all()

dp/cloud/python/magma/db_service/tests/unit/test_b0cad5321c88.py

@@ -0,0 +1,38 @@
+from magma.db_service.tests.alembic_testcase import AlembicTestCase
+
+
+class Testb0cad5321c88TestCase(AlembicTestCase):
+
+    def setUp(self) -> None:
+        super().setUp()
+        self.tables = [
+            'cbsd_states',
+            'domain_proxy_logs',
+            'grant_states',
+            'request_states',
+            'request_types',
+            'cbsds',
+            'active_mode_configs',
+            'channels',
+            'requests',
+            'grants',
+            'responses',
+        ]
+        self.up_revision = "b0cad5321c88"
+
+    def test_b0cad5321c88_upgrade(self):
+        # given / when
+        self.upgrade()
+
+        # then
+        self.then_tables_are(self.tables)
+
+    def test_b0cad5321c88_downgrade(self):
+        # given
+        self.upgrade()
+
+        # when
+        self.downgrade()
+
+        # then
+        self.then_tables_are()

dp/cloud/python/magma/test_runner/tests/test_active_mode_controller.py

@@ -20,7 +20,7 @@
 from dp.protos.enodebd_dp_pb2_grpc import DPServiceStub
 from magma.db_service.db_initialize import DBInitializer
 from magma.db_service.session_manager import SessionManager
-from magma.db_service.tests.db_testcase import DBTestCase
+from magma.db_service.tests.db_testcase import BaseDBTestCase
 from magma.test_runner.config import TestConfig
 from retrying import retry
 
@@ -32,7 +32,7 @@
 
 
 @pytest.mark.local
-class ActiveModeControllerTestCase(DBTestCase):
+class ActiveModeControllerTestCase(BaseDBTestCase):
     @classmethod
     def setUpClass(cls):
         wait_for_elastic_to_start()

dp/cloud/python/magma/test_runner/tests/test_dp_with_orc8r.py

@@ -27,7 +27,7 @@
 from dp.protos.enodebd_dp_pb2_grpc import DPServiceStub
 from magma.db_service.db_initialize import DBInitializer
 from magma.db_service.session_manager import SessionManager
-from magma.db_service.tests.db_testcase import DBTestCase
+from magma.db_service.tests.db_testcase import BaseDBTestCase
 from magma.test_runner.config import TestConfig
 from retrying import retry
 
@@ -45,7 +45,7 @@
 
 
 @pytest.mark.orc8r
-class DomainProxyOrc8rTestCase(DBTestCase):
+class DomainProxyOrc8rTestCase(BaseDBTestCase):
     @classmethod
     def setUpClass(cls):
         wait_for_elastic_to_start()