Add rollback for tests

qiita_core/util.py

@@ -83,12 +83,17 @@ def class_modifier(cls):
         # Now, we decorate the setup and teardown functions
         class DecoratedClass(cls):
             def setUp(self):
+                # Add one to context so we can rollback in tearDown
+                qdb.sql_connection.TRN._contexts_entered = 1
+
                 super(DecoratedClass, self).setUp()
                 self.conn_handler = qdb.sql_connection.SQLConnectionHandler()
 
-            @qdb.environment_manager.reset_test_database
             def tearDown(self):
                 super(DecoratedClass, self).tearDown()
+                # Rollback the changes and force out of context to run clean up
+                qdb.sql_connection.TRN.rollback()
+                qdb.sql_connection.TRN.__exit__(None, None, None)
 
         return DecoratedClass
     return class_modifier
@@ -98,8 +103,7 @@ def execute_as_transaction(func):
     """Decorator to make a method execute inside a transaction"""
     @wraps(func)
     def wrapper(*args, **kwargs):
-        from qiita_db.sql_connection import TRN
-        with TRN:
+        with qdb.sql_connection.TRN:
             return func(*args, **kwargs)
     return wrapper
 

qiita_db/metadata_template/test/test_prep_template.py

@@ -861,7 +861,7 @@ def test_create_error_cleanup(self):
             }
         metadata = pd.DataFrame.from_dict(metadata_dict, orient='index')
 
-        exp_id = qdb.util.get_count("qiita.prep_template") + 1
+        exp_id = qdb.util.get_next_id("qiita.prep_template")
 
         with self.assertRaises(ValueError):
             qdb.metadata_template.prep_template.PrepTemplate.create(
@@ -961,26 +961,26 @@ def _common_creation_checks(self, new_id, pt, fp_count):
 
     def test_create(self):
         """Creates a new PrepTemplate"""
-        fp_count = qdb.util.get_count('qiita.filepath')
-        new_id = qdb.util.get_count('qiita.prep_template') + 1
+        fp_count = qdb.util.get_next_id('qiita.filepath')
+        new_id = qdb.util.get_next_id('qiita.prep_template')
         pt = qdb.metadata_template.prep_template.PrepTemplate.create(
             self.metadata, self.test_study, self.data_type)
         self._common_creation_checks(new_id, pt, fp_count)
 
     def test_create_already_prefixed_samples(self):
         """Creates a new PrepTemplate"""
-        fp_count = qdb.util.get_count('qiita.filepath')
-        new_id = qdb.util.get_count('qiita.prep_template') + 1
+        fp_count = qdb.util.get_next_id('qiita.filepath')
+        new_id = qdb.util.get_next_id('qiita.prep_template')
         pt = npt.assert_warns(
             qdb.exceptions.QiitaDBWarning,
             qdb.metadata_template.prep_template.PrepTemplate.create,
             self.metadata_prefixed, self.test_study, self.data_type)
         self._common_creation_checks(new_id, pt, fp_count)
 
     def test_generate_files(self):
-        fp_count = qdb.util.get_count("qiita.filepath")
+        fp_count = qdb.util.get_next_id("qiita.filepath")
         self.tester.generate_files()
-        obs = qdb.util.get_count("qiita.filepath")
+        obs = qdb.util.get_next_id("qiita.filepath")
         # We just make sure that the count has been increased by 2, since
         # the contents of the files have been tested elsewhere.
         self.assertEqual(obs, fp_count + 2)
@@ -1007,17 +1007,17 @@ def test_create_qiime_mapping_file(self):
 
     def test_create_data_type_id(self):
         """Creates a new PrepTemplate passing the data_type_id"""
-        fp_count = qdb.util.get_count('qiita.filepath')
-        new_id = qdb.util.get_count('qiita.prep_template') + 1
+        fp_count = qdb.util.get_next_id('qiita.filepath')
+        new_id = qdb.util.get_next_id('qiita.prep_template')
         pt = qdb.metadata_template.prep_template.PrepTemplate.create(
             self.metadata, self.test_study, self.data_type_id)
         self._common_creation_checks(new_id, pt, fp_count)
 
     def test_create_warning(self):
         """Warns if a required columns is missing for a given functionality
         """
-        fp_count = qdb.util.get_count("qiita.filepath")
-        new_id = qdb.util.get_count('qiita.prep_template') + 1
+        fp_count = qdb.util.get_next_id("qiita.filepath")
+        new_id = qdb.util.get_next_id('qiita.prep_template')
         del self.metadata['barcode']
         pt = npt.assert_warns(
             qdb.exceptions.QiitaDBWarning,

qiita_db/test/test_analysis.py

@@ -129,7 +129,7 @@ def test_has_access_no_access(self):
     def test_create(self):
         sql = "SELECT EXTRACT(EPOCH FROM NOW())"
         time1 = float(self.conn_handler.execute_fetchall(sql)[0][0])
-        new_id = qdb.util.get_count("qiita.analysis") + 1
+        new_id = qdb.util.get_next_id("qiita.analysis")
         new = qdb.analysis.Analysis.create(
             qdb.user.User("admin@foo.bar"), "newAnalysis", "A New Analysis")
         self.assertEqual(new.id, new_id)
@@ -148,7 +148,7 @@ def test_create(self):
         self.assertEqual(obs, [[new_id, 1]])
 
     def test_create_nonqiita_portal(self):
-        new_id = qdb.util.get_count("qiita.analysis") + 1
+        new_id = qdb.util.get_next_id("qiita.analysis")
         qiita_config.portal = "EMP"
         qdb.analysis.Analysis.create(
             qdb.user.User("admin@foo.bar"), "newAnalysis", "A New Analysis")
@@ -162,7 +162,7 @@ def test_create_nonqiita_portal(self):
     def test_create_parent(self):
         sql = "SELECT EXTRACT(EPOCH FROM NOW())"
         time1 = float(self.conn_handler.execute_fetchall(sql)[0][0])
-        new_id = qdb.util.get_count("qiita.analysis") + 1
+        new_id = qdb.util.get_next_id("qiita.analysis")
         new = qdb.analysis.Analysis.create(
             qdb.user.User("admin@foo.bar"), "newAnalysis", "A New Analysis",
             qdb.analysis.Analysis(1))
@@ -180,7 +180,7 @@ def test_create_parent(self):
         self.assertEqual(obs, [[1, new_id]])
 
     def test_create_from_default(self):
-        new_id = qdb.util.get_count("qiita.analysis") + 1
+        new_id = qdb.util.get_next_id("qiita.analysis")
         owner = qdb.user.User("test@foo.bar")
         new = qdb.analysis.Analysis.create(
             owner, "newAnalysis", "A New Analysis", from_default=True)
@@ -206,23 +206,23 @@ def test_create_from_default(self):
     def test_exists(self):
         qiita_config.portal = 'QIITA'
         self.assertTrue(qdb.analysis.Analysis.exists(1))
-        new_id = qdb.util.get_count("qiita.analysis") + 1
+        new_id = qdb.util.get_next_id("qiita.analysis")
         self.assertFalse(qdb.analysis.Analysis.exists(new_id))
         qiita_config.portal = 'EMP'
         self.assertFalse(qdb.analysis.Analysis.exists(1))
-        new_id = qdb.util.get_count("qiita.analysis") + 1
+        new_id = qdb.util.get_next_id("qiita.analysis")
         self.assertFalse(qdb.analysis.Analysis.exists(new_id))
 
     def test_delete(self):
         # successful delete
-        total_analyses = qdb.util.get_count("qiita.analysis")
+        total_analyses = qdb.util.get_next_id("qiita.analysis")
         qdb.analysis.Analysis.delete(1)
         self.assertEqual(total_analyses - 1,
-                         qdb.util.get_count("qiita.analysis"))
+                         qdb.util.get_next_id("qiita.analysis"))
 
         # no possible to delete
         with self.assertRaises(qdb.exceptions.QiitaDBUnknownIDError):
-            qdb.analysis.Analysis.delete(total_analyses + 1)
+            qdb.analysis.Analysis.delete(total_analyses)
 
     def test_retrieve_owner(self):
         self.assertEqual(self.analysis.owner, qdb.user.User("test@foo.bar"))
@@ -379,7 +379,7 @@ def test_retrieve_biom_tables_empty(self):
         self.assertEqual(new.biom_tables, {})
 
     def test_set_step(self):
-        new_id = qdb.util.get_count("qiita.analysis") + 1
+        new_id = qdb.util.get_next_id("qiita.analysis")
         new = qdb.analysis.Analysis.create(
             qdb.user.User("admin@foo.bar"), "newAnalysis",
             "A New Analysis", qdb.analysis.Analysis(1))
@@ -389,7 +389,7 @@ def test_set_step(self):
         self.assertEqual(obs, [[new_id, 2]])
 
     def test_set_step_twice(self):
-        new_id = qdb.util.get_count("qiita.analysis") + 1
+        new_id = qdb.util.get_next_id("qiita.analysis")
         new = qdb.analysis.Analysis.create(
             qdb.user.User("admin@foo.bar"), "newAnalysis", "A New Analysis",
             qdb.analysis.Analysis(1))
@@ -513,7 +513,7 @@ def test_unshare(self):
         self.assertEqual(self.analysis.shared_with, [])
 
     def test_build_mapping_file(self):
-        new_id = qdb.util.get_count('qiita.filepath') + 1
+        new_id = qdb.util.get_next_id('qiita.filepath')
         samples = {4: ['1.SKB8.640193', '1.SKD8.640184', '1.SKB7.640196']}
 
         npt.assert_warns(qdb.exceptions.QiitaDBWarning,
@@ -571,7 +571,7 @@ def test_build_mapping_file_duplicated_samples_merge(self):
         assert_frame_equal(obs, exp)
 
     def test_build_biom_tables(self):
-        new_id = qdb.util.get_count('qiita.filepath') + 1
+        new_id = qdb.util.get_next_id('qiita.filepath')
         samples = {4: ['1.SKB8.640193', '1.SKD8.640184', '1.SKB7.640196']}
         self.analysis._build_biom_tables(samples, 100)
         obs = self.analysis.biom_tables
@@ -665,7 +665,7 @@ def test_build_files_raises_value_error(self):
             self.analysis.build_files(-10)
 
     def test_add_file(self):
-        new_id = qdb.util.get_count('qiita.filepath') + 1
+        new_id = qdb.util.get_next_id('qiita.filepath')
         fp = self.get_fp('testfile.txt')
         with open(fp, 'w') as f:
             f.write('testfile!')

qiita_db/test/test_job.py

@@ -221,7 +221,7 @@ def test_create_exists(self):
 
     def test_create_exists_return_existing(self):
         """Makes sure creation doesn't duplicate a job by returning existing"""
-        new_id = qdb.util.get_count("qiita.analysis") + 1
+        new_id = qdb.util.get_next_id("qiita.analysis")
         qdb.analysis.Analysis.create(qdb.user.User("demo@microbio.me"), "new",
                                      "desc")
         sql = """INSERT INTO qiita.analysis_sample
@@ -298,7 +298,7 @@ def test_retrieve_error_exists(self):
         self.assertEqual(self.job.error.msg, "TESTERROR")
 
     def test_add_results(self):
-        fp_count = qdb.util.get_count('qiita.filepath')
+        fp_count = qdb.util.get_next_id('qiita.filepath')
         self.job.add_results([(join(self._job_folder, "1_job_result.txt"),
                              "plain_text")])
 
@@ -309,7 +309,7 @@ def test_add_results(self):
         self.assertEqual(obs, [[1, 13], [1, fp_count + 1]])
 
     def test_add_results_dir(self):
-        fp_count = qdb.util.get_count('qiita.filepath')
+        fp_count = qdb.util.get_next_id('qiita.filepath')
         # Create a test directory
         test_dir = join(self._job_folder, "2_test_folder")
 

qiita_db/test/test_study.py

@@ -347,7 +347,7 @@ def test_create_duplicate(self):
     def test_create_study_min_data(self):
         """Insert a study into the database"""
         before = datetime.now()
-        new_id = qdb.util.get_count('qiita.study') + 1
+        new_id = qdb.util.get_next_id('qiita.study')
         obs = qdb.study.Study.create(
             qdb.user.User('test@foo.bar'), "Fried chicken microbiome", [1],
             self.info)

qiita_db/util.py

@@ -960,6 +960,19 @@ def get_count(table):
         return qdb.sql_connection.TRN.execute_fetchlast()
 
 
+def get_next_id(table):
+    """Returns the next ID in the series for the table
+
+       Returns
+       int
+           The next ID that will be assigned in the table
+    """
+    with qdb.sql_connection.TRN:
+        sql = "SELECT nextval({0}_{0}_id_seq) FROM {0}".format(table)
+        qdb.sql_connection.TRN.add(sql)
+        return qdb.sql_connection.TRN.execute_fetchlast()
+
+
 def check_count(table, exp_count):
     """Checks that the number of rows in a table equals the expected count