quay/test/test_storageproxy.py

import os
import requests
import unittest

from flask import Flask
from flask.ext.testing import LiveServerTestCase

from initdb import setup_database_for_testing, finished_database_for_testing
from storage import Storage
from util.security.instancekeys import InstanceKeys

_PORT_NUMBER = 5001

class TestStorageProxy(LiveServerTestCase):
  def setUp(self):
    setup_database_for_testing(self)

  def tearDown(self):
    finished_database_for_testing(self)

  def create_app(self):
    global _PORT_NUMBER
    _PORT_NUMBER = _PORT_NUMBER + 1

    self.test_app = Flask('teststorageproxy')
    self.test_app.config['LIVESERVER_PORT'] = _PORT_NUMBER

    if os.environ.get('DEBUG') == 'true':
      self.test_app.config['DEBUG'] = True

    self.test_app.config['TESTING'] = True
    self.test_app.config['SERVER_HOSTNAME'] = 'localhost:%s' % _PORT_NUMBER

    self.test_app.config['INSTANCE_SERVICE_KEY_KID_LOCATION'] = 'test/data/test.kid'
    self.test_app.config['INSTANCE_SERVICE_KEY_LOCATION'] = 'test/data/test.pem'
    self.test_app.config['INSTANCE_SERVICE_KEY_SERVICE'] = 'quay'

    # UGH... Such a stupid hack!
    self.test_app.config['FEATURE_PROXY_STORAGE'] = self.id().find('notinstalled') < 0

    self.test_app.config['DISTRIBUTED_STORAGE_CONFIG'] = {
      'test': ['FakeStorage', {}],
    }

    instance_keys = InstanceKeys(self.test_app)
    self.storage = Storage(self.test_app, instance_keys=instance_keys)
    self.test_app.config['DISTRIBUTED_STORAGE_PREFERENCE'] = ['test']
    return self.test_app

  def test_storage_proxy_auth_notinstalled(self):
    # Active direct download on the fake storage.
    self.storage.put_content(['test'], 'supports_direct_download', 'true')

    # Get the unwrapped URL.
    direct_download_url = self.storage.get_direct_download_url(['test'], 'somepath')
    self.assertEquals(-1, direct_download_url.find('/_storage_proxy/'))

    # Ensure that auth returns 404.
    headers = {
      'X-Original-URI': 'someurihere'
    }

    resp = requests.get('http://%s/_storage_proxy_auth' % self.test_app.config['SERVER_HOSTNAME'],
                        headers=headers)
    self.assertEquals(404, resp.status_code)


  def test_storage_proxy_auth(self):
    # Active direct download on the fake storage.
    self.storage.put_content(['test'], 'supports_direct_download', 'true')

    # Get the wrapped URL.
    direct_download_url = self.storage.get_direct_download_url(['test'], 'somepath')

    # Ensure it refers to the storage proxy.
    proxy_index = direct_download_url.find('/_storage_proxy/')
    self.assertTrue(proxy_index > 0)

    # Ensure that auth returns 200 for the URL pieces.
    headers = {
      'X-Original-URI': direct_download_url[proxy_index:]
    }

    resp = requests.get('http://%s/_storage_proxy_auth' % self.test_app.config['SERVER_HOSTNAME'],
                        headers=headers)
    self.assertEquals(200, resp.status_code)


if __name__ == '__main__':
  unittest.main()