import sqlalchemy as sa from test.lib import testing from test.lib.testing import adict from test.lib.engines import drop_all_tables import sys class TablesTest(testing.TestBase): """An integration test that creates and uses tables.""" # 'once', 'each', None run_setup_bind = 'once' # 'once', 'each', None run_define_tables = 'once' # 'once', 'each', None run_create_tables = 'once' # 'once', 'each', None run_inserts = 'each' # 'each', None run_deletes = 'each' # 'once', 'each', None run_dispose_bind = None bind = None metadata = None tables = None other = None @classmethod def setup_class(cls): cls._init_class() cls._setup_once_tables() cls._setup_once_inserts() @classmethod def _init_class(cls): if cls.run_define_tables == 'each': if cls.run_create_tables == 'once': cls.run_create_tables = 'each' assert cls.run_inserts in ('each', None) if cls.other is None: cls.other = adict() if cls.tables is None: cls.tables = adict() if cls.bind is None: setattr(cls, 'bind', cls.setup_bind()) if cls.metadata is None: setattr(cls, 'metadata', sa.MetaData()) if cls.metadata.bind is None: cls.metadata.bind = cls.bind @classmethod def _setup_once_inserts(cls): if cls.run_inserts == 'once': cls._load_fixtures() cls.insert_data() @classmethod def _setup_once_tables(cls): if cls.run_define_tables == 'once': cls.define_tables(cls.metadata) if cls.run_create_tables == 'once': cls.metadata.create_all() cls.tables.update(cls.metadata.tables) def _setup_each_tables(self): if self.run_define_tables == 'each': self.tables.clear() if self.run_create_tables == 'each': drop_all_tables(self.metadata) self.metadata.clear() self.define_tables(self.metadata) if self.run_create_tables == 'each': self.metadata.create_all() self.tables.update(self.metadata.tables) def _setup_each_inserts(self): if self.run_inserts == 'each': self._load_fixtures() self.insert_data() def _teardown_each_tables(self): # no need to run deletes if tables are recreated on setup if self.run_define_tables != 'each' and self.run_deletes: for table in reversed(self.metadata.sorted_tables): try: table.delete().execute().close() except sa.exc.DBAPIError, ex: print >> sys.stderr, "Error emptying table %s: %r" % ( table, ex) def _setup_each_bind(self): if self.setup_bind == 'each': setattr(cls, 'bind', self.setup_bind()) def _teardown_each_bind(self): if self.run_dispose_bind == 'each': self.dispose_bind(self.bind) def setup(self): self._setup_each_bind() self._setup_each_tables() self._setup_each_inserts() def teardown(self): self._teardown_each_tables() self._teardown_each_bind() @classmethod def _teardown_once_metadata_bind(cls): if cls.run_create_tables: cls.metadata.drop_all() if cls.run_dispose_bind == 'once': cls.dispose_bind(cls.bind) cls.metadata.bind = None if cls.run_setup_bind is not None: cls.bind = None @classmethod def teardown_class(cls): cls._teardown_once_metadata_bind() @classmethod def setup_bind(cls): return testing.db @classmethod def dispose_bind(cls, bind): if hasattr(bind, 'dispose'): bind.dispose() elif hasattr(bind, 'close'): bind.close() @classmethod def define_tables(cls, metadata): pass @classmethod def fixtures(cls): return {} @classmethod def insert_data(cls): pass def sql_count_(self, count, fn): self.assert_sql_count(self.bind, fn, count) def sql_eq_(self, callable_, statements, with_sequences=None): self.assert_sql(self.bind, callable_, statements, with_sequences) def _load_fixtures(self): headers, rows = {}, {} for table, data in self.fixtures().iteritems(): if isinstance(table, basestring): table = self.tables[table] headers[table] = data[0] rows[table] = data[1:] for table in self.metadata.sorted_tables: if table not in headers: continue table.bind.execute( table.insert(), [dict(zip(headers[table], column_values)) for column_values in rows[table]])