X-Git-Url: https://git.decadent.org.uk/gitweb/?a=blobdiff_plain;f=tests%2Fdb_test.py;h=d7f906cd34a147ad79916b283c0122be0761e34a;hb=4830d9be143c7645ea932b53fae095e275ad7814;hp=e9a533650df2e910ec4584b6990fca6d9158119e;hpb=3f70dafbc130cbebfe6b0a814831b05121411d4f;p=dak.git

diff --git a/tests/db_test.py b/tests/db_test.py
index e9a53365..d7f906cd 100644
--- a/tests/db_test.py
+++ b/tests/db_test.py
@@ -3,20 +3,47 @@ from base_test import DakTestCase, fixture
 from daklib.config import Config
 from daklib.dbconn import DBConn
 
-from sqlalchemy import create_engine
+from sqlalchemy import create_engine, __version__
 from sqlalchemy.exc import SADeprecationWarning
+from sqlalchemy.schema import DDL
 
 import pickle
 import warnings
 
-# suppress some deprecation warnings in squeeze related to sqlalchemy
-warnings.filterwarnings('ignore', \
-    "The SQLAlchemy PostgreSQL dialect has been renamed from 'postgres' to 'postgresql'.*", \
-    SADeprecationWarning)
+all_tables = ['architecture', 'archive', 'bin_associations', 'bin_contents',
+    'binaries', 'binary_acl', 'binary_acl_map', 'build_queue', 'build_queue_files',
+    'changes', 'changes_pending_binaries', 'changes_pending_files',
+    'changes_pending_files_map', 'changes_pending_source',
+    'changes_pending_source_files', 'changes_pool_files', 'component', 'config',
+    'dsc_files', 'files', 'fingerprint', 'keyring_acl_map', 'keyrings', 'location',
+    'maintainer', 'new_comments', 'override', 'override_type', 'policy_queue',
+    'priority', 'section', 'source', 'source_acl', 'src_associations',
+    'src_format', 'src_uploaders', 'suite', 'suite_architectures',
+    'suite_build_queue_copy', 'suite_src_formats', 'uid', 'upload_blocks']
+
+drop_plpgsql = "DROP LANGUAGE IF EXISTS plpgsql CASCADE"
+create_plpgsql = "CREATE LANGUAGE plpgsql"
+create_function = """CREATE OR REPLACE FUNCTION tfunc_set_modified() RETURNS trigger AS $$
+    BEGIN NEW.modified = now(); return NEW; END;
+    $$ LANGUAGE 'plpgsql'"""
+create_trigger = """CREATE TRIGGER modified_%s BEFORE UPDATE ON %s
+    FOR EACH ROW EXECUTE PROCEDURE tfunc_set_modified()"""
 
 class DBDakTestCase(DakTestCase):
+    def execute(self, statement):
+        DDL(statement).execute(self.metadata.bind)
+
+    def create_all_triggers(self):
+        for statement in (drop_plpgsql, create_plpgsql, create_function):
+            self.execute(statement)
+        for table in all_tables:
+            self.execute(create_trigger % (table, table))
+
     def setUp(self):
         cnf = Config()
+        if cnf["DB::Name"] in ('backports', 'obscurity', 'projectb'):
+            self.fail("You have configured an invalid database name: '%s'." % \
+                    cnf["DB::Name"])
         if cnf["DB::Host"]:
             # TCP/IP
             connstr = "postgres://%s" % cnf["DB::Host"]
@@ -29,16 +56,28 @@ class DBDakTestCase(DakTestCase):
             if cnf["DB::Port"] and cnf["DB::Port"] != "-1":
                 connstr += "?port=%s" % cnf["DB::Port"]
 
-        pickle_file = open(fixture('db-metadata.pkl'), 'r')
+        pickle_filename = 'db-metadata-%s.pkl' % __version__
+        pickle_file = open(fixture(pickle_filename), 'r')
         self.metadata = pickle.load(pickle_file)
         self.metadata.ddl_listeners = pickle.load(pickle_file)
         pickle_file.close()
         self.metadata.bind = create_engine(connstr)
         self.metadata.create_all()
+        self.create_all_triggers()
         self.session = DBConn().session()
 
+    def classes_to_clean(self):
+        """
+        The function classes_to_clean() returns a list of classes. All objects
+        of each class will be deleted from the database in tearDown(). This
+        function should be overridden in derived test cases as needed.
+        """
+        return ()
+
     def tearDown(self):
-        #pass
-        self.session.close()
+        self.session.rollback()
+        for class_ in self.classes_to_clean():
+            self.session.query(class_).delete()
+        self.session.commit()
         #self.metadata.drop_all()