cloudbridge/tests/test_download_driver.py

"""
Provider-agnostic unit tests for the base ranged download driver
(``BaseBucketObject.download_to_file`` / ``_download_ranged``).

The driver is the engine behind transparent large downloads on providers that
do not override it (GCP, OpenStack Swift). Because the mock provider is
AWS-backed and AWS overrides the driver with boto3's native downloader, the
driver is exercised here directly against in-memory fakes so it has coverage
in CI without cloud credentials.
"""
import os
import shutil
import tempfile
import threading
import unittest

from cloudbridge.base.resources import BaseBucketObject
from cloudbridge.interfaces.exceptions import InvalidValueException
from cloudbridge.interfaces.resources import TransferConfig


class _Recorder:
    """Thread-safe range log shared by the original and cloned fake
    services."""

    def __init__(self, content):
        self.content = content
        self._lock = threading.Lock()
        self.ranges = []            # (offset, length) served
        self.services_used = set()  # id() of each service that served a range
        self.clone_count = 0
        self.single_shot = False
        self.active = 0
        self.max_active = 0
        self.fail_on_offset = None  # offset that should raise
        self.on_serve = None        # hook called as each range is served

    def serve_range(self, service, offset, length):
        with self._lock:
            self.active += 1
            self.max_active = max(self.max_active, self.active)
        try:
            if self.on_serve:
                self.on_serve()
            if self.fail_on_offset == offset:
                raise RuntimeError("boom at offset %d" % offset)
            # Hold briefly so concurrent fetches genuinely overlap.
            threading.Event().wait(0.02)
            with self._lock:
                self.ranges.append((offset, length))
                self.services_used.add(id(service))
            return self.content[offset:offset + length]
        finally:
            with self._lock:
                self.active -= 1


class _FakeService:
    def __init__(self, recorder, provider):
        self._recorder = recorder
        self._provider = provider

    def download_range(self, bucket, object_name, offset, length):
        return self._recorder.serve_range(self, offset, length)


class _FakeStorage:
    def __init__(self, service):
        self._bucket_objects = service


class _FakeProvider:
    def __init__(self, recorder):
        self._recorder = recorder
        self.storage = _FakeStorage(_FakeService(recorder, self))

    def clone(self, zone=None):
        self._recorder.clone_count += 1
        return _FakeProvider(self._recorder)

    def _get_config_value(self, key, default_value=None):
        return default_value


class _DriverObject(BaseBucketObject):
    """A BaseBucketObject wired to fakes with tiny transfer sizes."""

    def __init__(self, provider, threshold, part_size, concurrency):
        super(_DriverObject, self).__init__(provider)
        self._threshold = threshold
        self._part_size = part_size
        self._concurrency = concurrency

    @property
    def id(self):
        return "obj"

    @property
    def name(self):
        return "obj"

    @property
    def size(self):
        return len(self._provider._recorder.content)

    @property
    def bucket(self):
        return "BUCKET"

    def save_content(self, target_stream):
        self._provider._recorder.single_shot = True
        target_stream.write(self._provider._recorder.content)

    def _multipart_threshold(self, config=None):
        if config is not None and config.threshold is not None:
            return config.threshold
        return self._threshold

    def _multipart_part_size(self, config=None):
        if config is not None and config.part_size is not None:
            return config.part_size
        return self._part_size

    def _multipart_max_concurrency(self, config=None):
        if config is not None and config.max_concurrency is not None:
            return config.max_concurrency
        return self._concurrency


class DownloadDriverTestCase(unittest.TestCase):

    def _driver(self, recorder, threshold, part_size, concurrency):
        return _DriverObject(
            _FakeProvider(recorder), threshold, part_size, concurrency)

    def _download(self, driver, config=None):
        fd, path = tempfile.mkstemp()
        os.close(fd)
        os.remove(path)
        try:
            driver.download_to_file(path, config)
            with open(path, 'rb') as f:
                return f.read()
        finally:
            if os.path.exists(path):
                os.remove(path)

    def test_reassembles_content_in_order(self):
        content = b"abcdefghijABCDEFGHIJ0123456789x"  # 31 bytes -> 8 ranges
        recorder = _Recorder(content)
        driver = self._driver(
            recorder, threshold=10, part_size=4, concurrency=3)
        self.assertEqual(self._download(driver), content)
        self.assertFalse(recorder.single_shot)
        # Ranges tile the object exactly: no gaps, no overlap, short tail.
        self.assertEqual(
            sorted(recorder.ranges),
            [(offset, min(4, 31 - offset)) for offset in range(0, 31, 4)])

    def test_below_threshold_uses_single_shot(self):
        content = b"tiny content"
        recorder = _Recorder(content)
        driver = self._driver(
            recorder, threshold=100, part_size=4, concurrency=3)
        self.assertEqual(self._download(driver), content)
        self.assertTrue(recorder.single_shot)
        self.assertEqual(recorder.ranges, [])

    def test_downloads_ranges_concurrently_via_cloned_services(self):
        concurrency = 4
        content = bytes(range(12))  # 12 ranges of one byte each
        recorder = _Recorder(content)
        driver = self._driver(
            recorder, threshold=1, part_size=1, concurrency=concurrency)
        self.assertEqual(self._download(driver), content)

        # A clone per worker, reused across ranges.
        self.assertEqual(recorder.clone_count, concurrency)
        self.assertEqual(len(recorder.services_used), concurrency)
        # Real parallelism happened, bounded by the configured concurrency.
        self.assertGreater(recorder.max_active, 1)
        self.assertLessEqual(recorder.max_active, concurrency)

    def test_single_concurrency_does_not_clone(self):
        content = b"abcdefghij"
        recorder = _Recorder(content)
        driver = self._driver(
            recorder, threshold=1, part_size=4, concurrency=1)
        self.assertEqual(self._download(driver), content)
        self.assertEqual(recorder.clone_count, 0)
        self.assertEqual(recorder.max_active, 1)

    def test_per_call_config_overrides_concurrency(self):
        content = bytes(range(12))
        recorder = _Recorder(content)
        driver = self._driver(
            recorder, threshold=1, part_size=1, concurrency=1)
        result = None
        fd, path = tempfile.mkstemp()
        os.close(fd)
        try:
            driver.download_to_file(path, TransferConfig(max_concurrency=3))
            with open(path, 'rb') as f:
                result = f.read()
        finally:
            os.remove(path)
        self.assertEqual(result, content)
        self.assertEqual(recorder.clone_count, 3)
        self.assertGreater(recorder.max_active, 1)
        self.assertLessEqual(recorder.max_active, 3)

    def test_removes_partial_file_and_raises_on_range_failure(self):
        content = bytes(range(16))
        recorder = _Recorder(content)
        recorder.fail_on_offset = 8
        driver = self._driver(
            recorder, threshold=1, part_size=4, concurrency=2)
        fd, path = tempfile.mkstemp()
        os.close(fd)
        os.remove(path)
        try:
            with self.assertRaises(Exception):
                driver.download_to_file(path)
            self.assertFalse(os.path.exists(path))
        finally:
            if os.path.exists(path):
                os.remove(path)

    def test_destination_only_appears_once_complete(self):
        content = bytes(range(256))
        recorder = _Recorder(content)
        driver = self._driver(
            recorder, threshold=1, part_size=16, concurrency=3)
        fd, path = tempfile.mkstemp()
        os.close(fd)
        os.remove(path)
        seen_early = []
        recorder.on_serve = lambda: seen_early.append(os.path.exists(path))
        try:
            driver.download_to_file(path)
            with open(path, 'rb') as f:
                self.assertEqual(f.read(), content)
        finally:
            if os.path.exists(path):
                os.remove(path)
        # A partially written object is never visible at the destination.
        self.assertTrue(seen_early)
        self.assertNotIn(True, seen_early)

    def test_survives_concurrent_downloader_taking_the_destination(self):
        # Galaxy gives every download of a dataset the same cache .tmp path,
        # so a second download of the same dataset can rename the destination
        # away while this one is still fetching ranges.
        content = bytes(range(256))
        recorder = _Recorder(content)
        driver = self._driver(
            recorder, threshold=1, part_size=16, concurrency=3)
        directory = tempfile.mkdtemp()
        path = os.path.join(directory, 'dataset.dat')
        taken = os.path.join(directory, 'taken.dat')

        def steal_destination():
            if os.path.exists(path):
                os.replace(path, taken)

        recorder.on_serve = steal_destination
        try:
            driver.download_to_file(path)
            with open(path, 'rb') as f:
                self.assertEqual(f.read(), content)
        finally:
            shutil.rmtree(directory)

    def test_failed_download_leaves_an_existing_destination_intact(self):
        content = bytes(range(64))
        recorder = _Recorder(content)
        recorder.fail_on_offset = 16
        driver = self._driver(
            recorder, threshold=1, part_size=16, concurrency=2)
        directory = tempfile.mkdtemp()
        path = os.path.join(directory, 'dataset.dat')
        with open(path, 'wb') as f:
            f.write(b'previously cached')
        try:
            with self.assertRaises(Exception):
                driver.download_to_file(path)
            # The cached copy survives a failed refetch, and no scratch file
            # is left behind next to it.
            with open(path, 'rb') as f:
                self.assertEqual(f.read(), b'previously cached')
            self.assertEqual(os.listdir(directory), ['dataset.dat'])
        finally:
            shutil.rmtree(directory)

    def test_part_size_must_be_positive(self):
        content = bytes(range(16))
        recorder = _Recorder(content)
        driver = self._driver(
            recorder, threshold=1, part_size=4, concurrency=2)
        with self.assertRaises(InvalidValueException):
            self._download(driver, TransferConfig(part_size=0))
        self.assertEqual(recorder.ranges, [])


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