68 lines
2.7 KiB
Python
68 lines
2.7 KiB
Python
import pytest
|
|
|
|
from contextlib import contextmanager
|
|
from mock import patch
|
|
from moto import mock_s3
|
|
import boto
|
|
|
|
from app import config_provider
|
|
from storage import CloudFrontedS3Storage, StorageContext
|
|
from util.ipresolver import IPResolver
|
|
from util.ipresolver.test.test_ipresolver import test_aws_ip, aws_ip_range_data
|
|
from test.fixtures import *
|
|
|
|
_TEST_CONTENT = os.urandom(1024)
|
|
_TEST_BUCKET = 'some_bucket'
|
|
_TEST_USER = 'someuser'
|
|
_TEST_PASSWORD = 'somepassword'
|
|
_TEST_PATH = 'some/cool/path'
|
|
|
|
@pytest.fixture(params=[True, False])
|
|
def ipranges_populated(request):
|
|
return request.param
|
|
|
|
@mock_s3
|
|
def test_direct_download(test_aws_ip, aws_ip_range_data, ipranges_populated, app):
|
|
ipresolver = IPResolver(app)
|
|
if ipranges_populated:
|
|
empty_range_data = {
|
|
'syncToken': 123456789,
|
|
'prefixes': [],
|
|
}
|
|
|
|
with patch.object(ipresolver, '_get_aws_ip_ranges', lambda: aws_ip_range_data if ipranges_populated else empty_range_data):
|
|
context = StorageContext('nyc', None, None, config_provider, ipresolver)
|
|
|
|
# Create a test bucket and put some test content.
|
|
boto.connect_s3().create_bucket(_TEST_BUCKET)
|
|
|
|
engine = CloudFrontedS3Storage(context, 'cloudfrontdomain', 'keyid', 'test/data/test.pem', 'some/path',
|
|
_TEST_BUCKET, _TEST_USER, _TEST_PASSWORD)
|
|
engine.put_content(_TEST_PATH, _TEST_CONTENT)
|
|
assert engine.exists(_TEST_PATH)
|
|
|
|
# Request a direct download URL for a request from a known AWS IP, and ensure we are returned an S3 URL.
|
|
assert 's3.amazonaws.com' in engine.get_direct_download_url(_TEST_PATH, test_aws_ip)
|
|
|
|
if ipranges_populated:
|
|
# Request a direct download URL for a request from a non-AWS IP, and ensure we are returned a CloudFront URL.
|
|
assert 'cloudfrontdomain' in engine.get_direct_download_url(_TEST_PATH, '1.2.3.4')
|
|
else:
|
|
# Request a direct download URL for a request from a non-AWS IP, but since IP Ranges isn't populated, we still
|
|
# get back an S3 URL.
|
|
assert 's3.amazonaws.com' in engine.get_direct_download_url(_TEST_PATH, '1.2.3.4')
|
|
|
|
@mock_s3
|
|
def test_direct_download_no_ip(test_aws_ip, aws_ip_range_data, ipranges_populated, app):
|
|
ipresolver = IPResolver(app)
|
|
context = StorageContext('nyc', None, None, config_provider, ipresolver)
|
|
|
|
# Create a test bucket and put some test content.
|
|
boto.connect_s3().create_bucket(_TEST_BUCKET)
|
|
|
|
engine = CloudFrontedS3Storage(context, 'cloudfrontdomain', 'keyid', 'test/data/test.pem', 'some/path',
|
|
_TEST_BUCKET, _TEST_USER, _TEST_PASSWORD)
|
|
engine.put_content(_TEST_PATH, _TEST_CONTENT)
|
|
assert engine.exists(_TEST_PATH)
|
|
|
|
assert 'cloudfrontdomain' in engine.get_direct_download_url(_TEST_PATH)
|