Merge remote-tracking branch 'upstream/master' into python-registry-v2

test/registry_tests.py

@@ -7,12 +7,14 @@ from flask.blueprints import Blueprint
 from flask.ext.testing import LiveServerTestCase
 
 from app import app
+from data.database import close_db_filter, configure
 from endpoints.v1 import v1_bp
 from endpoints.v2 import v2_bp
 from endpoints.v2.manifest import SignedManifestBuilder
 from endpoints.api import api_bp
 from initdb import wipe_database, initialize_database, populate_database
 from endpoints.csrf import generate_csrf_token
+from tempfile import NamedTemporaryFile
 
 import endpoints.decorated
 import json
@@ -20,6 +22,7 @@ import features
 import hashlib
 
 import tarfile
+import shutil
 
 from jwkest.jws import SIGNER_ALGS
 from jwkest.jwk import RSAKey
@@ -37,7 +40,9 @@ except ValueError:
   pass
 
 
-# Add a test blueprint for generating CSRF tokens and setting feature flags.
+# Add a test blueprint for generating CSRF tokens, setting feature flags and reloading the
+# DB connection.
+
 testbp = Blueprint('testbp', __name__)
 
 @testbp.route('/csrf', methods=['GET'])
@@ -51,6 +56,15 @@ def set_feature(feature_name):
   features._FEATURES[feature_name].value = request.get_json()['value']
   return jsonify({'old_value': old_value})
 
+@testbp.route('/reloaddb', methods=['POST'])
+def reload_db():
+  # Close any existing connection.
+  close_db_filter(None)
+
+  # Reload the database config.
+  configure(app.config)
+  return 'OK'
+
 app.register_blueprint(testbp, url_prefix='/__test')
 
 
@@ -78,6 +92,62 @@ class TestFeature(object):
                            headers={'Content-Type': 'application/json'})
 
 
+_PORT_NUMBER = 5001
+_CLEAN_DATABASE_PATH = None
+
+
+def get_new_database_uri():
+  # If a clean copy of the database has not yet been created, create one now.
+  global _CLEAN_DATABASE_PATH
+  if not _CLEAN_DATABASE_PATH:
+    wipe_database()
+    initialize_database()
+    populate_database()
+    close_db_filter(None)
+
+    # Save the path of the clean database.
+    _CLEAN_DATABASE_PATH = app.config['TEST_DB_FILE'].name
+
+  # Create a new temp file to be used as the actual backing database for the test.
+  # Note that we have the close() the file to ensure we can copy to it via shutil.
+  local_db_file = NamedTemporaryFile(delete=True)
+  local_db_file.close()
+
+  # Copy the clean database to the path.
+  shutil.copy2(_CLEAN_DATABASE_PATH, local_db_file.name)
+  return 'sqlite:///{0}'.format(local_db_file.name)
+
+
+class RegistryTestCase(LiveServerTestCase):
+  maxDiff = None
+
+  def create_app(self):
+    global _PORT_NUMBER
+    _PORT_NUMBER = _PORT_NUMBER + 1
+    app.config['TESTING'] = True
+    app.config['LIVESERVER_PORT'] = _PORT_NUMBER
+    app.config['DB_URI'] = get_new_database_uri()
+    return app
+
+  def setUp(self):
+    self.clearSession()
+
+    # Tell the remote running app to reload the database. By default, the app forks from the
+    # current context and has already loaded the DB config with the *original* DB URL. We call
+    # the remote reload method to force it to pick up the changes to DB_URI set in the create_app
+    # method.
+    self.conduct('POST', '/__test/reloaddb')
+
+  def clearSession(self):
+    self.session = requests.Session()
+    self.signature = None
+    self.docker_token = 'true'
+
+    # Load the CSRF token.
+    self.csrf_token = ''
+    self.csrf_token = self.conduct('GET', '/__test/csrf').text
+
+
 class BaseRegistryMixin(object):
   def conduct(self, method, url, headers=None, data=None, auth=None, params=None, expected_code=200):
     params = params or {}