clean-python/tests/api_client/test_sync_files.py

import io
from unittest import mock

import pytest
from urllib3.response import HTTPResponse
from urllib3.util.request import set_file_position

from clean_python.api_client import ApiException
from clean_python.api_client import download_file
from clean_python.api_client import download_fileobj
from clean_python.api_client import upload_file
from clean_python.api_client import upload_fileobj
from clean_python.api_client.files import _SeekableChunkIterator
from clean_python.api_client.files import DEFAULT_UPLOAD_TIMEOUT

MODULE = "clean_python.api_client.files"


@pytest.fixture
def pool():
    pool = mock.Mock()
    return pool


@pytest.fixture
def responses_single():
    return [
        HTTPResponse(
            body=b"X" * 42,
            headers={"Content-Range": "bytes 0-41/42"},
            status=206,
        )
    ]


@pytest.fixture
def responses_double():
    return [
        HTTPResponse(
            body=b"X" * 64,
            headers={"Content-Range": "bytes 0-63/65"},
            status=206,
        ),
        HTTPResponse(
            body=b"X",
            headers={"Content-Range": "bytes 64-64/65"},
            status=206,
        ),
    ]


def test_download_fileobj(pool, responses_single):
    stream = io.BytesIO()
    pool.request.side_effect = responses_single
    download_fileobj("some-url", stream, chunk_size=64, pool=pool)

    pool.request.assert_called_with(
        "GET",
        "some-url",
        headers={"Range": "bytes=0-63"},
        timeout=5.0,
    )
    assert stream.tell() == 42


def test_download_fileobj_two_chunks(pool, responses_double):
    stream = io.BytesIO()
    pool.request.side_effect = responses_double

    callback_func = mock.Mock()

    download_fileobj(
        "some-url", stream, chunk_size=64, pool=pool, callback_func=callback_func
    )

    (_, kwargs1), (_, kwargs2) = pool.request.call_args_list
    assert kwargs1["headers"] == {"Range": "bytes=0-63"}
    assert kwargs2["headers"] == {"Range": "bytes=64-127"}
    assert stream.tell() == 65

    # Check callback func
    (args1, _), (args2, _) = callback_func.call_args_list

    assert args1 == (63, 65)
    assert args2 == (65, 65)


def test_download_fileobj_no_multipart(pool, responses_single):
    """The remote server does not support range requests"""
    responses_single[0].status = 200
    pool.request.side_effect = responses_single

    with pytest.raises(ApiException) as e:
        download_fileobj("some-url", None, chunk_size=64, pool=pool)

    assert e.value.status == 200
    assert str(e.value) == "200: The file server does not support multipart downloads."


def test_download_fileobj_forbidden(pool, responses_single):
    """The remote server does not support range requests"""
    responses_single[0].status = 403
    pool.request.side_effect = responses_single

    with pytest.raises(ApiException) as e:
        download_fileobj("some-url", None, chunk_size=64, pool=pool)

    assert e.value.status == 403


@mock.patch(MODULE + ".download_fileobj")
def test_download_file(download_fileobj, tmp_path):
    download_file(
        "http://domain/a.b",
        tmp_path / "c.d",
        chunk_size=64,
        timeout=3.0,
        pool="foo",
        headers_factory="bar",
    )

    args, kwargs = download_fileobj.call_args
    assert args[0] == "http://domain/a.b"
    assert isinstance(args[1], io.IOBase)
    assert args[1].mode == "wb"
    assert args[1].name == str(tmp_path / "c.d")
    assert kwargs["chunk_size"] == 64
    assert kwargs["timeout"] == 3.0
    assert kwargs["pool"] == "foo"
    assert kwargs["headers_factory"] == "bar"


@mock.patch(MODULE + ".download_fileobj")
def test_download_file_directory(download_fileobj, tmp_path):
    # if a target directory is specified, a filename is created from the url
    download_file("http://domain/a.b", tmp_path, chunk_size=64, timeout=3.0, pool="foo")

    args, kwargs = download_fileobj.call_args
    assert args[1].name == str(tmp_path / "a.b")


@pytest.fixture
def upload_response():
    return HTTPResponse(status=200)


@pytest.fixture
def fileobj():
    stream = io.BytesIO()
    stream.write(b"X" * 39)
    stream.seek(0)
    return stream


@pytest.mark.parametrize(
    "chunk_size,expected_body",
    [
        (64, [b"X" * 39]),
        (39, [b"X" * 39]),
        (38, [b"X" * 38, b"X"]),
        (16, [b"X" * 16, b"X" * 16, b"X" * 7]),
    ],
)
def test_upload_fileobj(pool, fileobj, upload_response, chunk_size, expected_body):
    pool.request.return_value = upload_response
    upload_fileobj("some-url", fileobj, chunk_size=chunk_size, pool=pool)

    args, kwargs = pool.request.call_args
    assert args == ("PUT", "some-url")
    assert list(kwargs["body"]) == expected_body
    assert kwargs["headers"] == {"Content-Length": "39"}
    assert kwargs["timeout"] == DEFAULT_UPLOAD_TIMEOUT


def test_upload_fileobj_callback(pool, fileobj, upload_response):
    expected_body = [b"X" * 16, b"X" * 16, b"X" * 7]
    chunk_size = 16

    pool.request.return_value = upload_response
    callback_func = mock.Mock()

    upload_fileobj(
        "some-url",
        fileobj,
        chunk_size=chunk_size,
        pool=pool,
        callback_func=callback_func,
    )

    args, kwargs = pool.request.call_args
    assert args == ("PUT", "some-url")
    assert list(kwargs["body"]) == expected_body
    assert kwargs["headers"] == {"Content-Length": "39"}
    assert kwargs["timeout"] == DEFAULT_UPLOAD_TIMEOUT

    # Check callback_func
    (args1, _), (args2, _), (args3, _) = callback_func.call_args_list
    assert args1 == (16, 39)
    assert args2 == (32, 39)
    assert args3 == (39, 39)


def test_upload_fileobj_with_md5(pool, fileobj, upload_response):
    pool.request.return_value = upload_response
    upload_fileobj("some-url", fileobj, pool=pool, md5=b"abcd")

    # base64.b64encode(b"abcd")).decode()
    expected_md5 = "YWJjZA=="

    args, kwargs = pool.request.call_args
    assert kwargs["headers"] == {"Content-Length": "39", "Content-MD5": expected_md5}


def test_upload_fileobj_empty_file():
    with pytest.raises(IOError, match="The file object is empty."):
        upload_fileobj("some-url", io.BytesIO())


def test_upload_fileobj_non_binary_file():
    with pytest.raises(IOError, match="The file object is not in binary*"):
        upload_fileobj("some-url", io.StringIO())


def test_upload_fileobj_errors(pool, fileobj, upload_response):
    upload_response.status = 400
    pool.request.return_value = upload_response
    with pytest.raises(ApiException) as e:
        upload_fileobj("some-url", fileobj, pool=pool)

    assert e.value.status == 400


@mock.patch(MODULE + ".upload_fileobj")
def test_upload_file(upload_fileobj, tmp_path):
    path = tmp_path / "myfile"
    with path.open("wb") as f:
        f.write(b"X")

    upload_file(
        "http://domain/a.b",
        path,
        chunk_size=1234,
        timeout=3.0,
        pool="foo",
        md5=b"abcd",
        headers_factory="bar",
    )

    args, kwargs = upload_fileobj.call_args
    assert args[0] == "http://domain/a.b"
    assert isinstance(args[1], io.IOBase)
    assert args[1].mode == "rb"
    assert args[1].name == str(path)
    assert kwargs["timeout"] == 3.0
    assert kwargs["chunk_size"] == 1234
    assert kwargs["pool"] == "foo"
    assert kwargs["md5"] == b"abcd"
    assert kwargs["headers_factory"] == "bar"


def test_seekable_chunk_iterator():
    data = b"XYZ"
    body = _SeekableChunkIterator(io.BytesIO(data), chunk_size=4)

    pos = set_file_position(body, pos=0)
    assert pos == 0
    assert list(body) == [data]
    assert list(body) == []
    set_file_position(body, pos)
    assert list(body) == [data]


def test_download_fileobj_with_headers(pool, responses_single):
    pool.request.side_effect = responses_single
    download_fileobj(
        "some-url",
        io.BytesIO(),
        chunk_size=64,
        pool=pool,
        headers_factory=lambda: {"foo": "bar"},
    )

    pool.request.assert_called_with(
        "GET",
        "some-url",
        headers={"Range": "bytes=0-63", "foo": "bar"},
        timeout=5.0,
    )


def test_upload_fileobj_with_headers(pool, fileobj, upload_response):
    pool.request.return_value = upload_response
    upload_fileobj(
        "some-url", fileobj, pool=pool, headers_factory=lambda: {"foo": "bar"}
    )

    _, kwargs = pool.request.call_args
    assert kwargs["headers"] == {"Content-Length": "39", "foo": "bar"}


def test_upload_fileobj_201_response(pool, fileobj):
    pool.request.return_value = HTTPResponse(status=201)
    # no error is raised
    upload_fileobj(
        "some-url", fileobj, pool=pool, headers_factory=lambda: {"foo": "bar"}
    )

    pool.request.assert_called_once()