quay/storage/test/test_cloud_storage.py

import os

from StringIO import StringIO

import pytest

import moto
import boto

from storage import S3Storage, StorageContext
from storage.cloud import _CloudStorage, _PartUploadMetadata
from storage.cloud import _CHUNKS_KEY

_TEST_CONTENT = os.urandom(1024)
_TEST_BUCKET = 'some_bucket'
_TEST_USER = 'someuser'
_TEST_PASSWORD = 'somepassword'
_TEST_PATH = 'some/cool/path'
_TEST_CONTEXT = StorageContext('nyc', None, None, None, None)

@pytest.fixture()
def storage_engine():
  mock = moto.mock_s3()
  mock.start()

  # Create a test bucket and put some test content.
  boto.connect_s3().create_bucket(_TEST_BUCKET)
  engine = S3Storage(_TEST_CONTEXT, 'some/path', _TEST_BUCKET, _TEST_USER, _TEST_PASSWORD)
  engine.put_content(_TEST_PATH, _TEST_CONTENT)

  yield engine

  mock.stop()


def test_basicop(storage_engine):
  # Ensure the content exists.
  assert storage_engine.exists(_TEST_PATH)

  # Verify it can be retrieved.
  assert storage_engine.get_content(_TEST_PATH) == _TEST_CONTENT

  # Retrieve a checksum for the content.
  storage_engine.get_checksum(_TEST_PATH)

  # Remove the file.
  storage_engine.remove(_TEST_PATH)

  # Ensure it no longer exists.
  with pytest.raises(IOError):
    storage_engine.get_content(_TEST_PATH)

  with pytest.raises(IOError):
    storage_engine.get_checksum(_TEST_PATH)

  assert not storage_engine.exists(_TEST_PATH)


def test_copy_samecreds(storage_engine):
  # Copy the content to another engine.
  another_engine = S3Storage(_TEST_CONTEXT, 'another/path', _TEST_BUCKET, _TEST_USER,
                             _TEST_PASSWORD)
  storage_engine.copy_to(another_engine, _TEST_PATH)

  # Verify it can be retrieved.
  assert another_engine.get_content(_TEST_PATH) == _TEST_CONTENT


def test_copy_differentcreds(storage_engine):
  # Copy the content to another engine.
  another_engine = S3Storage(_TEST_CONTEXT, 'another/path', 'another_bucket', 'blech',
                             'password')
  boto.connect_s3().create_bucket('another_bucket')

  storage_engine.copy_to(another_engine, _TEST_PATH)

  # Verify it can be retrieved.
  assert another_engine.get_content(_TEST_PATH) == _TEST_CONTENT


def test_stream_read(storage_engine):
  # Read the streaming content.
  data = ''.join(storage_engine.stream_read(_TEST_PATH))
  assert data == _TEST_CONTENT


def test_stream_read_file(storage_engine):
  with storage_engine.stream_read_file(_TEST_PATH) as f:
    assert f.read() == _TEST_CONTENT


def test_stream_write(storage_engine):
  new_data = os.urandom(4096)
  storage_engine.stream_write(_TEST_PATH, StringIO(new_data), content_type='Cool/Type')
  assert storage_engine.get_content(_TEST_PATH) == new_data


@pytest.mark.parametrize('chunk_count', [
  0,
  1,
  50,
])
@pytest.mark.parametrize('force_client_side', [
  False,
  True
])
def test_chunk_upload(storage_engine, chunk_count, force_client_side):
  if chunk_count == 0 and force_client_side:
    return
  
  upload_id, metadata = storage_engine.initiate_chunked_upload()
  final_data = ''

  for index in range(0, chunk_count):
    chunk_data = os.urandom(1024)
    final_data = final_data + chunk_data
    bytes_written, new_metadata, error = storage_engine.stream_upload_chunk(upload_id, 0,
                                                                            len(chunk_data),
                                                                            StringIO(chunk_data),
                                                                            metadata)
    metadata = new_metadata

    assert bytes_written == len(chunk_data)
    assert error is None
    assert len(metadata[_CHUNKS_KEY]) == index + 1

  # Complete the chunked upload.
  storage_engine.complete_chunked_upload(upload_id, 'some/chunked/path', metadata,
                                         force_client_side=force_client_side)

  # Ensure the file contents are valid.
  assert storage_engine.get_content('some/chunked/path') == final_data


@pytest.mark.parametrize('chunk_count', [
  0,
  1,
  50,
])
def test_cancel_chunked_upload(storage_engine, chunk_count):
  upload_id, metadata = storage_engine.initiate_chunked_upload()

  for _ in range(0, chunk_count):
    chunk_data = os.urandom(1024)
    _, new_metadata, _ = storage_engine.stream_upload_chunk(upload_id, 0,
                                                            len(chunk_data),
                                                            StringIO(chunk_data),
                                                            metadata)
    metadata = new_metadata

  # Cancel the upload.
  storage_engine.cancel_chunked_upload(upload_id, metadata)

  # Ensure all chunks were deleted.
  for chunk in metadata[_CHUNKS_KEY]:
    assert not storage_engine.exists(chunk.path)


def test_large_chunks_upload(storage_engine):
  # Make the max chunk size much smaller for testing.
  storage_engine.maximum_chunk_size = storage_engine.minimum_chunk_size * 2

  upload_id, metadata = storage_engine.initiate_chunked_upload()

  # Write a "super large" chunk, to ensure that it is broken into smaller chunks.
  chunk_data = os.urandom(int(storage_engine.maximum_chunk_size * 2.5))
  bytes_written, new_metadata, _ = storage_engine.stream_upload_chunk(upload_id, 0,
                                                                      -1,
                                                                      StringIO(chunk_data),
                                                                      metadata)
  assert len(chunk_data) == bytes_written

  # Complete the chunked upload.
  storage_engine.complete_chunked_upload(upload_id, 'some/chunked/path', new_metadata)

  # Ensure the file contents are valid.
  assert len(chunk_data) == len(storage_engine.get_content('some/chunked/path'))
  assert storage_engine.get_content('some/chunked/path') == chunk_data


def test_large_chunks_with_ragged_edge(storage_engine):
  # Make the max chunk size much smaller for testing and force it to have a ragged edge.
  storage_engine.maximum_chunk_size = storage_engine.minimum_chunk_size * 2 + 10

  upload_id, metadata = storage_engine.initiate_chunked_upload()

  # Write a few "super large" chunks, to ensure that it is broken into smaller chunks.
  all_data = ''
  for _ in range(0, 2):
    chunk_data = os.urandom(int(storage_engine.maximum_chunk_size) + 20)
    bytes_written, new_metadata, _ = storage_engine.stream_upload_chunk(upload_id, 0,
                                                                        -1,
                                                                        StringIO(chunk_data),
                                                                        metadata)
    assert len(chunk_data) == bytes_written
    all_data = all_data + chunk_data
    metadata = new_metadata

  # Complete the chunked upload.
  storage_engine.complete_chunked_upload(upload_id, 'some/chunked/path', new_metadata)

  # Ensure the file contents are valid.
  assert len(all_data) == len(storage_engine.get_content('some/chunked/path'))
  assert storage_engine.get_content('some/chunked/path') == all_data


@pytest.mark.parametrize('max_size, parts', [
  (50, [
    _PartUploadMetadata('foo', 0, 50),
    _PartUploadMetadata('foo', 50, 50),
  ]),

  (40, [
    _PartUploadMetadata('foo', 0, 25),
    _PartUploadMetadata('foo', 25, 25),
    _PartUploadMetadata('foo', 50, 25),
    _PartUploadMetadata('foo', 75, 25)
  ]),

  (51, [
    _PartUploadMetadata('foo', 0, 50),
    _PartUploadMetadata('foo', 50, 50),
  ]),

  (49, [
    _PartUploadMetadata('foo', 0, 25),
    _PartUploadMetadata('foo', 25, 25),
    _PartUploadMetadata('foo', 50, 25),
    _PartUploadMetadata('foo', 75, 25),
  ]),

  (99, [
    _PartUploadMetadata('foo', 0, 50),
    _PartUploadMetadata('foo', 50, 50),
  ]),

  (100, [
    _PartUploadMetadata('foo', 0, 100),
  ]),
])
def test_rechunked(max_size, parts):
  chunk = _PartUploadMetadata('foo', 0, 100)
  rechunked = list(_CloudStorage._rechunk(chunk, max_size))
  assert len(rechunked) == len(parts)
  for index, chunk in enumerate(rechunked):
    assert chunk == parts[index]