Postgre test fixtures

test_postgis_base.py

UNIT_TEST_DB_CONFIG = {'user': 'unittest', 'password': 'unittest',
                       'database': 'unittest', 'host': 'localhost', 'port': 5432}
                       
JDBC_URL = 'postgresql://{user}:{password}@{host}:{port}/{database}'.format(**UNIT_TEST_DB_CONFIG)

class SchemaManager:
    EXTRA_DROP_ACTIONS = []

    def __init__(self, engine: sqa_engine, schema_base=AnalysisSchemaBase):
        self.engine = engine
        self.schema_base = schema_base

    def create_tables(self):
        self.schema_base.metadata.create_all(self.engine)

    def drop_tables(self):
        table_list = ",".join(self.schema_base.metadata.tables.keys())
        self.engine.execute("drop table {} cascade".format(table_list))
        for action in self.EXTRA_DROP_ACTIONS:
            self.engine.execute(action)

    def __enter__(self):
        self.create_tables()

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.drop_tables()


class PostgisSchemaTestBaseClass(unittest.TestCase):
    def setUp(cls):
        cls.engine = create_engine(JDBC_URL)  # type: engine
        cls.schema_manager = SchemaManager(cls.engine)


class PostgisTestBaseClass(PostgisSchemaTestBaseClass):
    def setUp(self):
        super(PostgisTestBaseClass, self).setUp()
        self.schema_manager.create_tables()
        self.sessionmaker = sessionmaker(bind=self.engine, autocommit=False)

    def tearDown(self):
        super(PostgisTestBaseClass, self).tearDown()
        self.schema_manager.drop_tables()

    @contextmanager
    def get_session(self):
        """Provide a transactional scope around a series of operations."""
        session = self.sessionmaker()
        try:
            yield session
        except:
            session.rollback()
            raise
        finally:
            session.close()

    def hydrate_from_sql_file(self, filepath: str):
        with open(filepath, 'r') as sql_fp:
            sql = sql_fp.readlines()
        for line in sql:
            if line.startswith('--') or line.strip() == '':
                continue
            try:
                self.engine.execute(line)
            except Exception as ex:
                print(str(ex))

test_sample.py

class MySampleTestClass(PostgisTestBaseClass):
      def setUp(self):
        super(MySampleTestClass, self).setUp()
        self.hydrate_from_sql_file(SQL_FILE) # file path to SQL commands to hydrate data into database
        
      def run_a_simple_test(self):
        with self.get_session() as sesh:
          some_data = sesh.query(SomeStuff).all()
          self.assertTrue(...) # some tests here