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