Make testing much faster by using a save point, rather than recreating the database every test

initdb.py

@@ -29,7 +29,6 @@ SAMPLE_CMDS = [["/bin/bash"],
 REFERENCE_DATE = datetime(2013, 6, 23)
 TEST_STRIPE_ID = 'cus_2tmnh3PkXQS8NG'
 
-
 def __gen_checksum(image_id):
   h = hashlib.md5(image_id)
   return 'tarsum+sha256:' + h.hexdigest() + h.hexdigest()
@@ -112,6 +111,45 @@ def __generate_repository(user, name, description, is_public, permissions,
 
   return repo
 
+from util.dbutil import savepoint_sqlite
+
+db_initialized_for_testing = False
+testcases = {}
+
+def finished_database_for_testing(testcase):
+  """ Called when a testcase has finished using the database, indicating that
+      any changes should be discarded.
+  """
+  global testcases
+  testcases[testcase]['savepoint'].__exit__(True, None, None)
+
+def setup_database_for_testing(testcase):
+  """ Called when a testcase has started using the database, indicating that
+      the database should be setup (if not already) and a savepoint created.
+  """
+
+  # Sanity check to make sure we're not killing our prod db
+  db = model.db
+  if (not isinstance(model.db, SqliteDatabase) or
+      app.config['DB_DRIVER'] is not SqliteDatabase):
+    raise RuntimeError('Attempted to wipe production database!')
+
+  global db_initialized_for_testing
+  if not db_initialized_for_testing:
+    logger.debug('Setting up DB for testing.') 
+
+    # Setup the database.
+    wipe_database()
+    initialize_database()
+    populate_database()
+
+    db_initialized_for_testing = True
+
+  # Create a savepoint for the testcase.
+  global testcases  
+  testcases[testcase] = {}
+  testcases[testcase]['savepoint'] = savepoint_sqlite(db)
+  testcases[testcase]['savepoint'].__enter__()
 
 def initialize_database():
   create_model_tables(all_models)