test_files.py

import io
from unittest import mock

import pytest
from urllib3.response import HTTPResponse

from threedi_api_client.openapi import ApiException
from threedi_api_client.files import download_file
from threedi_api_client.files import download_fileobj
from threedi_api_client.files import upload_file
from threedi_api_client.files import upload_fileobj
from threedi_api_client.files import DEFAULT_UPLOAD_TIMEOUT


@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 e.value.reason == "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("threedi_api_client.files.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"
    )

    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"


@mock.patch("threedi_api_client.files.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("threedi_api_client.files.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"
    )

    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"