-
Notifications
You must be signed in to change notification settings - Fork 590
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
feat(dp): migration tests framework + poc test
Signed-off-by: Wojciech Sadowy <wojciech.sadowy@freedomfi.com>
- Loading branch information
Wojciech Sadowy
committed
Apr 28, 2022
1 parent
ca545de
commit a05c6b6
Showing
9 changed files
with
177 additions
and
36 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
75 changes: 75 additions & 0 deletions
75
dp/cloud/python/magma/db_service/tests/alembic_testcase.py
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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 assert_tables_equal(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())) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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() |
16 changes: 8 additions & 8 deletions
16
dp/cloud/python/magma/db_service/tests/local_db_test_case.py
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1,21 +1,21 @@ | ||
import testing.postgresql | ||
from magma.db_service.tests.db_testcase import DBTestCase | ||
|
||
Postgresql = testing.postgresql.PostgresqlFactory(cache_initialized_db=True) | ||
from magma.db_service.models import Base | ||
from magma.db_service.tests.db_testcase import BaseDBTestCase | ||
|
||
|
||
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() |
38 changes: 38 additions & 0 deletions
38
dp/cloud/python/magma/db_service/tests/unit/test_b0cad5321c88.py
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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.assert_tables_equal(self.tables) | ||
|
||
def test_b0cad5321c88_downgrade(self): | ||
# given | ||
self.upgrade() | ||
|
||
# when | ||
self.downgrade() | ||
|
||
# then | ||
self.assert_tables_equal() |
3 changes: 2 additions & 1 deletion
3
dp/cloud/python/magma/db_service/tests/unit/test_db_initialization.py
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters