<s> import io <EOL> import os <EOL> import alembic <EOL> from alembic import command <EOL> from alembic import config as alembic_config <EOL> from alembic import migration <EOL> from alembic import script as alembic_script <EOL> from oslo_config import cfg <EOL> from oslo_db . sqlalchemy import test_migrations as t_m <EOL> from oslo_log import log as logging <EOL> import sahara . db . migration <EOL> from sahara . db . sqlalchemy import api as sa <EOL> from sahara . db . sqlalchemy import model_base <EOL> from sahara . i18n import _LE <EOL> LOG = logging . getLogger ( __name__ ) <EOL> CONF = cfg . CONF <EOL> class BaseWalkMigrationTestCase ( object ) : <EOL> ALEMBIC_CONFIG = alembic_config . Config ( <EOL> os . path . join ( os . path . dirname ( sahara . db . migration . __file__ ) , <EOL> '<STR_LIT>' ) <EOL> ) <EOL> ALEMBIC_CONFIG . sahara_config = CONF <EOL> def _configure ( self , engine ) : <EOL> """<STR_LIT>""" <EOL> CONF . set_override ( '<STR_LIT>' , str ( engine . url ) , group = '<STR_LIT>' , <EOL> enforce_type = True ) <EOL> sa . cleanup ( ) <EOL> def _alembic_command ( self , alembic_command , engine , * args , ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> self . ALEMBIC_CONFIG . stdout = buf = io . StringIO ( ) <EOL> CONF . set_override ( '<STR_LIT>' , str ( engine . url ) , group = '<STR_LIT>' , <EOL> enforce_type = True ) <EOL> sa . cleanup ( ) <EOL> getattr ( command , alembic_command ) ( * args , ** kwargs ) <EOL> res = buf . getvalue ( ) . strip ( ) <EOL> LOG . debug ( '<STR_LIT>' . format ( <EOL> command = alembic_command , result = res ) ) <EOL> sa . cleanup ( ) <EOL> return res <EOL> def _get_versions ( self ) : <EOL> """<STR_LIT>""" <EOL> env = alembic_script . ScriptDirectory . from_config ( self . ALEMBIC_CONFIG ) <EOL> versions = [ ] <EOL> for rev in env . walk_revisions ( ) : <EOL> versions . append ( rev . revision ) <EOL> versions . reverse ( ) <EOL> return versions <EOL> def walk_versions ( self , engine = None ) : <EOL> self . _configure ( engine ) <EOL> versions = self . _get_versions ( ) <EOL> for ver in versions : <EOL> self . _migrate_up ( engine , ver , with_data = True ) <EOL> def _get_version_from_db ( self , engine ) : <EOL> """<STR_LIT>""" <EOL> conn = engine . connect ( ) <EOL> try : <EOL> context = migration . MigrationContext . configure ( conn ) <EOL> version = context . get_current_revision ( ) or '<STR_LIT>' <EOL> finally : <EOL> conn . close ( ) <EOL> return version <EOL> def _migrate ( self , engine , version , cmd ) : <EOL> """<STR_LIT>""" <EOL> self . _alembic_command ( cmd , engine , self . ALEMBIC_CONFIG , version ) <EOL> def _migrate_up ( self , engine , version , with_data = False ) : <EOL> """<STR_LIT>""" <EOL> check_version = version <EOL> try : <EOL> if with_data : <EOL> data = None <EOL> pre_upgrade = getattr ( <EOL> self , "<STR_LIT>" % check_version , None ) <EOL> if pre_upgrade : <EOL> data = pre_upgrade ( engine ) <EOL> self . _migrate ( engine , version , '<STR_LIT>' ) <EOL> self . assertEqual ( version , self . _get_version_from_db ( engine ) ) <EOL> if with_data : <EOL> check = getattr ( self , "<STR_LIT>" % check_version , None ) <EOL> if check : <EOL> check ( engine , data ) <EOL> except Exception : <EOL> LOG . error ( _LE ( "<STR_LIT>" <EOL> "<STR_LIT>" ) . format ( version = version , engine = engine ) ) <EOL> raise <EOL> class TestModelsMigrationsSync ( t_m . ModelsMigrationsSync ) : <EOL> """<STR_LIT>""" <EOL> ALEMBIC_CONFIG = alembic_config . Config ( <EOL> os . path . join ( os . path . dirname ( sahara . db . migration . __file__ ) , <EOL> '<STR_LIT>' ) <EOL> ) <EOL> ALEMBIC_CONFIG . sahara_config = CONF <EOL> def get_engine ( self ) : <EOL> return self . engine <EOL> def db_sync ( self , engine ) : <EOL> CONF . set_override ( '<STR_LIT>' , str ( engine . url ) , group = '<STR_LIT>' , <EOL> enforce_type = True ) <EOL> alembic . command . upgrade ( self . ALEMBIC_CONFIG , '<STR_LIT>' ) <EOL> def get_metadata ( self ) : <EOL> return model_base . SaharaBase . metadata </s>
