tests/db_test.py

   1 from base_test import DakTestCase, fixture
   2
   3 from daklib.config import Config
   4 from daklib.dbconn import DBConn
   5
   6 from sqlalchemy import create_engine, __version__
   7 from sqlalchemy.exc import SADeprecationWarning
   8 from sqlalchemy.schema import DDL
   9
  10 import pickle
  11 import warnings
  12
  13 # suppress some deprecation warnings in squeeze related to sqlalchemy
  14 warnings.filterwarnings('ignore', \
  15     "The SQLAlchemy PostgreSQL dialect has been renamed from 'postgres' to 'postgresql'.*", \
  16     SADeprecationWarning)
  17
  18 all_tables = ['architecture', 'archive', 'bin_associations', 'bin_contents',
  19     'binaries', 'binary_acl', 'binary_acl_map', 'build_queue', 'build_queue_files',
  20     'changes', 'changes_pending_binaries', 'changes_pending_files',
  21     'changes_pending_files_map', 'changes_pending_source',
  22     'changes_pending_source_files', 'changes_pool_files', 'component', 'config',
  23     'dsc_files', 'files', 'fingerprint', 'keyring_acl_map', 'keyrings', 'location',
  24     'maintainer', 'new_comments', 'override', 'override_type', 'policy_queue',
  25     'priority', 'section', 'source', 'source_acl', 'src_associations',
  26     'src_format', 'src_uploaders', 'suite', 'suite_architectures',
  27     'suite_build_queue_copy', 'suite_src_formats', 'uid', 'upload_blocks']
  28
  29 drop_plpgsql = "DROP LANGUAGE IF EXISTS plpgsql CASCADE"
  30 create_plpgsql = "CREATE LANGUAGE plpgsql"
  31 create_function = """CREATE OR REPLACE FUNCTION tfunc_set_modified() RETURNS trigger AS $$
  32     BEGIN NEW.modified = now(); return NEW; END;
  33     $$ LANGUAGE 'plpgsql'"""
  34 create_trigger = """CREATE TRIGGER modified_%s BEFORE UPDATE ON %s
  35     FOR EACH ROW EXECUTE PROCEDURE tfunc_set_modified()"""
  36
  37 class DBDakTestCase(DakTestCase):
  38     def execute(self, statement):
  39         DDL(statement).execute(self.metadata.bind)
  40
  41     def create_all_triggers(self):
  42         for statement in (drop_plpgsql, create_plpgsql, create_function):
  43             self.execute(statement)
  44         for table in all_tables:
  45             self.execute(create_trigger % (table, table))
  46
  47     def setUp(self):
  48         cnf = Config()
  49         if cnf["DB::Name"] in ('backports', 'obscurity', 'projectb'):
  50             self.fail("You have configured an invalid database name: '%s'." % \
  51                     cnf["DB::Name"])
  52         if cnf["DB::Host"]:
  53             # TCP/IP
  54             connstr = "postgres://%s" % cnf["DB::Host"]
  55             if cnf["DB::Port"] and cnf["DB::Port"] != "-1":
  56                 connstr += ":%s" % cnf["DB::Port"]
  57             connstr += "/%s" % cnf["DB::Name"]
  58         else:
  59             # Unix Socket
  60             connstr = "postgres:///%s" % cnf["DB::Name"]
  61             if cnf["DB::Port"] and cnf["DB::Port"] != "-1":
  62                 connstr += "?port=%s" % cnf["DB::Port"]
  63
  64         pickle_filename = 'db-metadata-%s.pkl' % __version__
  65         pickle_file = open(fixture(pickle_filename), 'r')
  66         self.metadata = pickle.load(pickle_file)
  67         self.metadata.ddl_listeners = pickle.load(pickle_file)
  68         pickle_file.close()
  69         self.metadata.bind = create_engine(connstr)
  70         self.metadata.create_all()
  71         self.create_all_triggers()
  72         self.session = DBConn().session()
  73
  74     def classes_to_clean(self):
  75         """
  76         The function classes_to_clean() returns a list of classes. All objects
  77         of each class will be deleted from the database in tearDown(). This
  78         function should be overridden in derived test cases as needed.
  79         """
  80         return ()
  81
  82     def tearDown(self):
  83         self.session.rollback()
  84         for class_ in self.classes_to_clean():
  85             self.session.query(class_).delete()
  86         self.session.commit()
  87         #self.metadata.drop_all()
  88