9194ef5abd
Docs/Test Workflow / Test docs build (push) Failing after 0s
Check links & references / links-check (push) Failing after 1s
Pytest/Test Workflow / Import Test and Pytest Run (ubuntu-latest, 3.10) (push) Failing after 0s
Pytest/Test Workflow / Import Test and Pytest Run (ubuntu-latest, 3.11) (push) Failing after 0s
PR Conflict Labeler / main (push) Failing after 2s
Pytest/Test Workflow / Import Test and Pytest Run (ubuntu-latest, 3.12) (push) Failing after 2s
Pytest/Test Workflow / Import Test and Pytest Run (ubuntu-latest, 3.13) (push) Failing after 0s
Pytest/Test Workflow / Build this Package (push) Failing after 5s
Pytest/Test Workflow / Import Test and Pytest Run (macos-latest, 3.10) (push) Has been cancelled
Pytest/Test Workflow / Import Test and Pytest Run (macos-latest, 3.11) (push) Has been cancelled
Pytest/Test Workflow / Import Test and Pytest Run (macos-latest, 3.12) (push) Has been cancelled
Pytest/Test Workflow / Import Test and Pytest Run (macos-latest, 3.13) (push) Has been cancelled
Pytest/Test Workflow / Import Test and Pytest Run (windows-latest, 3.10) (push) Has been cancelled
Pytest/Test Workflow / Import Test and Pytest Run (windows-latest, 3.11) (push) Has been cancelled
Pytest/Test Workflow / Import Test and Pytest Run (windows-latest, 3.12) (push) Has been cancelled
Pytest/Test Workflow / Import Test and Pytest Run (windows-latest, 3.13) (push) Has been cancelled
Pytest/Test Workflow / testing-guardian (push) Has been cancelled
420 lines
15 KiB
Python
420 lines
15 KiB
Python
from unittest.mock import MagicMock, mock_open, patch
|
|
|
|
import pytest
|
|
|
|
from supervision.assets.downloader import (
|
|
_download_asset,
|
|
download_assets,
|
|
is_md5_hash_matching,
|
|
)
|
|
from supervision.assets.list import ImageAssets, VideoAssets
|
|
|
|
|
|
class TestMD5HashMatching:
|
|
def test_file_exists_matching_hash(self) -> None:
|
|
"""Test is_md5_hash_matching when file exists and hash matches."""
|
|
test_content = b"test content"
|
|
test_hash = "9473fdd0d880a43c21b7778d34872157" # MD5 of "test content"
|
|
|
|
with (
|
|
patch("builtins.open", mock_open(read_data=test_content)),
|
|
patch("os.path.exists", return_value=True),
|
|
):
|
|
assert is_md5_hash_matching("dummy_file", test_hash)
|
|
|
|
def test_file_exists_not_matching_hash(self) -> None:
|
|
"""Test is_md5_hash_matching when file exists but hash doesn't match."""
|
|
test_content = b"test content"
|
|
wrong_hash = "wrong_hash"
|
|
|
|
with (
|
|
patch("builtins.open", mock_open(read_data=test_content)),
|
|
patch("os.path.exists", return_value=True),
|
|
):
|
|
assert not is_md5_hash_matching("dummy_file", wrong_hash)
|
|
|
|
def test_file_not_exists(self) -> None:
|
|
"""Test is_md5_hash_matching when file doesn't exist."""
|
|
with patch("os.path.exists", return_value=False):
|
|
assert not is_md5_hash_matching("nonexistent_file", "some_hash")
|
|
|
|
|
|
class TestDownloadAssets:
|
|
@patch("os.replace")
|
|
@patch("supervision.assets.downloader.logger")
|
|
@patch("supervision.assets.downloader.is_md5_hash_matching", return_value=True)
|
|
@patch("pathlib.Path.exists", return_value=True)
|
|
def test_already_exists_and_valid(
|
|
self, mock_exists, mock_md5, mock_logger, mock_replace
|
|
) -> None:
|
|
"""Test download_assets when file already exists and is valid."""
|
|
filename = "vehicles.mp4"
|
|
result = download_assets(filename)
|
|
assert result == filename
|
|
mock_logger.info.assert_called_with("%s asset download complete.", filename)
|
|
|
|
@patch("os.replace")
|
|
@patch("supervision.assets.downloader.logger")
|
|
@patch("os.remove")
|
|
@patch(
|
|
"supervision.assets.downloader.is_md5_hash_matching",
|
|
side_effect=[False, True],
|
|
)
|
|
@patch("pathlib.Path.open", new_callable=mock_open)
|
|
@patch("pathlib.Path.mkdir")
|
|
@patch("pathlib.Path.exists", return_value=True)
|
|
@patch("supervision.assets.downloader.copyfileobj")
|
|
@patch("supervision.assets.downloader.tqdm")
|
|
@patch("supervision.assets.downloader.get")
|
|
def test_already_exists_but_corrupted(
|
|
self,
|
|
mock_get,
|
|
mock_tqdm,
|
|
mock_copyfileobj,
|
|
mock_exists,
|
|
mock_mkdir,
|
|
mock_open_file,
|
|
mock_md5,
|
|
mock_remove,
|
|
mock_logger,
|
|
mock_replace,
|
|
) -> None:
|
|
"""Test download_assets when file exists but is corrupted (re-downloads)."""
|
|
filename = "vehicles.mp4"
|
|
|
|
mock_response = MagicMock()
|
|
mock_response.headers = {"Content-Length": "100"}
|
|
mock_response.raw = MagicMock()
|
|
mock_response.raise_for_status = MagicMock()
|
|
mock_get.return_value = mock_response
|
|
|
|
mock_tqdm.wrapattr.return_value.__enter__ = MagicMock(
|
|
return_value=mock_response.raw
|
|
)
|
|
mock_tqdm.wrapattr.return_value.__exit__ = MagicMock(return_value=False)
|
|
|
|
result = download_assets(filename)
|
|
|
|
assert result == filename
|
|
mock_get.assert_called_once()
|
|
mock_copyfileobj.assert_called_once()
|
|
mock_logger.warning.assert_called_once_with("File corrupted. Re-downloading...")
|
|
mock_remove.assert_called_once_with(filename)
|
|
|
|
@patch("os.replace")
|
|
@patch("supervision.assets.downloader.logger")
|
|
@patch("supervision.assets.downloader.is_md5_hash_matching", return_value=True)
|
|
@patch("pathlib.Path.open", new_callable=mock_open)
|
|
@patch("pathlib.Path.mkdir")
|
|
@patch("pathlib.Path.exists", return_value=False)
|
|
@patch("supervision.assets.downloader.copyfileobj")
|
|
@patch("supervision.assets.downloader.tqdm")
|
|
@patch("supervision.assets.downloader.get")
|
|
def test_download_new_file(
|
|
self,
|
|
mock_get,
|
|
mock_tqdm,
|
|
mock_copyfileobj,
|
|
mock_exists,
|
|
mock_mkdir,
|
|
mock_open_file,
|
|
mock_md5,
|
|
mock_logger,
|
|
mock_replace,
|
|
) -> None:
|
|
"""Test download_assets verifies a freshly downloaded file."""
|
|
filename = "vehicles.mp4"
|
|
|
|
mock_response = MagicMock()
|
|
mock_response.headers = {"Content-Length": "100"}
|
|
mock_response.raw = MagicMock()
|
|
mock_response.raise_for_status = MagicMock()
|
|
mock_get.return_value = mock_response
|
|
|
|
mock_tqdm.wrapattr.return_value.__enter__ = MagicMock(
|
|
return_value=mock_response.raw
|
|
)
|
|
mock_tqdm.wrapattr.return_value.__exit__ = MagicMock(return_value=False)
|
|
|
|
result = download_assets(filename)
|
|
assert result == filename
|
|
mock_logger.info.assert_called_with("Downloading %s assets", filename)
|
|
mock_get.assert_called_once()
|
|
mock_response.raise_for_status.assert_called_once_with()
|
|
mock_copyfileobj.assert_called_once()
|
|
mock_md5.assert_called_once_with(filename, "8155ff4e4de08cfa25f39de96483f918")
|
|
|
|
@patch("os.replace")
|
|
@patch("supervision.assets.downloader.logger")
|
|
@patch("os.remove")
|
|
@patch(
|
|
"supervision.assets.downloader.is_md5_hash_matching",
|
|
side_effect=[False, True],
|
|
)
|
|
@patch("pathlib.Path.open", new_callable=mock_open)
|
|
@patch("pathlib.Path.mkdir")
|
|
@patch("pathlib.Path.exists", return_value=False)
|
|
@patch("supervision.assets.downloader.copyfileobj")
|
|
@patch("supervision.assets.downloader.tqdm")
|
|
@patch("supervision.assets.downloader.get")
|
|
def test_download_new_file_retries_corrupted_payload(
|
|
self,
|
|
mock_get,
|
|
mock_tqdm,
|
|
mock_copyfileobj,
|
|
mock_exists,
|
|
mock_mkdir,
|
|
mock_open_file,
|
|
mock_md5,
|
|
mock_remove,
|
|
mock_logger,
|
|
mock_replace,
|
|
) -> None:
|
|
"""Test download_assets retries once when a fresh payload fails MD5."""
|
|
filename = "vehicles.mp4"
|
|
|
|
mock_response = MagicMock()
|
|
mock_response.headers = {"Content-Length": "100"}
|
|
mock_response.raw = MagicMock()
|
|
mock_response.raise_for_status = MagicMock()
|
|
mock_get.return_value = mock_response
|
|
|
|
mock_tqdm.wrapattr.return_value.__enter__ = MagicMock(
|
|
return_value=mock_response.raw
|
|
)
|
|
mock_tqdm.wrapattr.return_value.__exit__ = MagicMock(return_value=False)
|
|
|
|
result = download_assets(filename)
|
|
|
|
assert result == filename
|
|
assert mock_get.call_count == 2
|
|
assert mock_copyfileobj.call_count == 2
|
|
mock_remove.assert_called_once_with(filename)
|
|
mock_logger.warning.assert_called_once_with("File corrupted. Re-downloading...")
|
|
|
|
@patch("os.replace")
|
|
@patch("supervision.assets.downloader.logger")
|
|
@patch("os.remove")
|
|
@patch(
|
|
"supervision.assets.downloader.is_md5_hash_matching",
|
|
side_effect=[False, False],
|
|
)
|
|
@patch("pathlib.Path.open", new_callable=mock_open)
|
|
@patch("pathlib.Path.mkdir")
|
|
@patch("pathlib.Path.exists", return_value=False)
|
|
@patch("supervision.assets.downloader.copyfileobj")
|
|
@patch("supervision.assets.downloader.tqdm")
|
|
@patch("supervision.assets.downloader.get")
|
|
def test_download_new_file_raises_after_second_md5_mismatch(
|
|
self,
|
|
mock_get,
|
|
mock_tqdm,
|
|
mock_copyfileobj,
|
|
mock_exists,
|
|
mock_mkdir,
|
|
mock_open_file,
|
|
mock_md5,
|
|
mock_remove,
|
|
mock_logger,
|
|
mock_replace,
|
|
) -> None:
|
|
"""Test download_assets fails after the verified retry is also corrupted."""
|
|
filename = "vehicles.mp4"
|
|
|
|
mock_response = MagicMock()
|
|
mock_response.headers = {"Content-Length": "100"}
|
|
mock_response.raw = MagicMock()
|
|
mock_response.raise_for_status = MagicMock()
|
|
mock_get.return_value = mock_response
|
|
|
|
mock_tqdm.wrapattr.return_value.__enter__ = MagicMock(
|
|
return_value=mock_response.raw
|
|
)
|
|
mock_tqdm.wrapattr.return_value.__exit__ = MagicMock(return_value=False)
|
|
|
|
with pytest.raises(ValueError, match="failed MD5 verification"):
|
|
download_assets(filename)
|
|
|
|
assert mock_get.call_count == 2
|
|
assert mock_copyfileobj.call_count == 2
|
|
assert mock_remove.call_count == 2
|
|
assert mock_logger.warning.call_count == 2
|
|
|
|
@patch("supervision.assets.downloader.logger")
|
|
@patch("supervision.assets.downloader.is_md5_hash_matching", return_value=True)
|
|
@patch("supervision.assets.downloader.copyfileobj")
|
|
@patch("supervision.assets.downloader.tqdm")
|
|
@patch("supervision.assets.downloader.get")
|
|
def test_download_new_file_to_custom_directory(
|
|
self,
|
|
mock_get,
|
|
mock_tqdm,
|
|
mock_copyfileobj,
|
|
mock_md5,
|
|
mock_logger,
|
|
tmp_path,
|
|
) -> None:
|
|
"""Test download_assets writes into an explicit output directory."""
|
|
filename = "vehicles.mp4"
|
|
target_directory = tmp_path / "nested" / "assets"
|
|
|
|
mock_response = MagicMock()
|
|
mock_response.headers = {"Content-Length": "100"}
|
|
mock_response.raw = MagicMock()
|
|
mock_response.raise_for_status = MagicMock()
|
|
mock_get.return_value = mock_response
|
|
|
|
mock_tqdm.wrapattr.return_value.__enter__ = MagicMock(
|
|
return_value=mock_response.raw
|
|
)
|
|
mock_tqdm.wrapattr.return_value.__exit__ = MagicMock(return_value=False)
|
|
mock_copyfileobj.side_effect = lambda _src, file: file.write(b"asset-bytes")
|
|
|
|
result = download_assets(filename, directory=target_directory)
|
|
|
|
assert result == str(target_directory / filename)
|
|
assert (target_directory / filename).exists()
|
|
assert (target_directory / filename).read_bytes() == b"asset-bytes"
|
|
mock_md5.assert_called_once_with(
|
|
str(target_directory / filename), "8155ff4e4de08cfa25f39de96483f918"
|
|
)
|
|
|
|
@patch("os.replace")
|
|
@patch("supervision.assets.downloader.get")
|
|
@patch("supervision.assets.downloader.tqdm")
|
|
@patch("supervision.assets.downloader.copyfileobj")
|
|
def test_partial_download_does_not_leave_final_file(
|
|
self,
|
|
mock_copyfileobj,
|
|
mock_tqdm,
|
|
mock_get,
|
|
mock_replace,
|
|
tmp_path,
|
|
) -> None:
|
|
"""Test _download_asset stages downloads so failed replaces do not leak."""
|
|
filename = "vehicles.mp4"
|
|
destination = tmp_path / filename
|
|
|
|
mock_response = MagicMock()
|
|
mock_response.headers = {"Content-Length": "100"}
|
|
mock_response.raw = MagicMock()
|
|
mock_response.raise_for_status = MagicMock()
|
|
mock_get.return_value = mock_response
|
|
|
|
mock_tqdm.wrapattr.return_value.__enter__ = MagicMock(
|
|
return_value=mock_response.raw
|
|
)
|
|
mock_tqdm.wrapattr.return_value.__exit__ = MagicMock(return_value=False)
|
|
|
|
mock_copyfileobj.side_effect = lambda _src, file: file.write(b"partial")
|
|
mock_replace.side_effect = OSError("boom")
|
|
|
|
with pytest.raises(OSError, match="boom"):
|
|
_download_asset(filename, destination)
|
|
|
|
mock_copyfileobj.assert_called_once()
|
|
assert not destination.exists()
|
|
assert not destination.with_name(f"{filename}.part").exists()
|
|
|
|
@patch("pathlib.Path.exists", return_value=False)
|
|
def test_invalid_asset(self, mock_exists) -> None:
|
|
"""Test download_assets with invalid asset name."""
|
|
invalid_filename = "invalid.mp4"
|
|
|
|
with pytest.raises(ValueError, match="Invalid asset") as exc_info:
|
|
download_assets(invalid_filename)
|
|
|
|
assert "Invalid asset" in str(exc_info.value)
|
|
assert "vehicles.mp4" in str(exc_info.value)
|
|
|
|
@patch("pathlib.Path.exists", return_value=True)
|
|
def test_invalid_asset_when_file_exists(self, mock_exists) -> None:
|
|
"""Test download_assets with invalid asset name that already exists."""
|
|
invalid_filename = "invalid.mp4"
|
|
|
|
with pytest.raises(ValueError, match="Invalid asset") as exc_info:
|
|
download_assets(invalid_filename)
|
|
|
|
assert "Invalid asset" in str(exc_info.value)
|
|
assert "vehicles.mp4" in str(exc_info.value)
|
|
|
|
@patch("os.replace")
|
|
@patch("supervision.assets.downloader.logger")
|
|
@patch("supervision.assets.downloader.is_md5_hash_matching", return_value=True)
|
|
@patch("pathlib.Path.open", new_callable=mock_open)
|
|
@patch("pathlib.Path.mkdir")
|
|
@patch("supervision.assets.downloader.copyfileobj")
|
|
@patch("supervision.assets.downloader.tqdm")
|
|
@patch("supervision.assets.downloader.get")
|
|
@patch("pathlib.Path.exists", return_value=False)
|
|
def test_with_video_enum(
|
|
self,
|
|
mock_exists,
|
|
mock_get,
|
|
mock_tqdm,
|
|
mock_copyfileobj,
|
|
mock_mkdir,
|
|
mock_open_file,
|
|
mock_md5,
|
|
mock_logger,
|
|
mock_replace,
|
|
) -> None:
|
|
"""Test download_assets with VideoAssets enum."""
|
|
asset = VideoAssets.VEHICLES
|
|
|
|
mock_response = MagicMock()
|
|
mock_response.headers = {"Content-Length": "100"}
|
|
mock_response.raw = MagicMock()
|
|
mock_response.raise_for_status = MagicMock()
|
|
mock_get.return_value = mock_response
|
|
|
|
mock_tqdm.wrapattr.return_value.__enter__ = MagicMock()
|
|
mock_tqdm.wrapattr.return_value.__exit__ = MagicMock()
|
|
|
|
result = download_assets(asset)
|
|
assert result == asset.filename
|
|
mock_logger.info.assert_called_with("Downloading %s assets", asset.filename)
|
|
mock_md5.assert_called_once_with(
|
|
asset.filename, "8155ff4e4de08cfa25f39de96483f918"
|
|
)
|
|
|
|
@patch("os.replace")
|
|
@patch("supervision.assets.downloader.logger")
|
|
@patch("supervision.assets.downloader.is_md5_hash_matching", return_value=True)
|
|
@patch("pathlib.Path.open", new_callable=mock_open)
|
|
@patch("pathlib.Path.mkdir")
|
|
@patch("supervision.assets.downloader.copyfileobj")
|
|
@patch("supervision.assets.downloader.tqdm")
|
|
@patch("supervision.assets.downloader.get")
|
|
@patch("pathlib.Path.exists", return_value=False)
|
|
def test_with_image_enum(
|
|
self,
|
|
mock_exists,
|
|
mock_get,
|
|
mock_tqdm,
|
|
mock_copyfileobj,
|
|
mock_mkdir,
|
|
mock_open_file,
|
|
mock_md5,
|
|
mock_logger,
|
|
mock_replace,
|
|
) -> None:
|
|
"""Test download_assets with ImageAssets enum."""
|
|
asset = ImageAssets.SOCCER
|
|
|
|
mock_response = MagicMock()
|
|
mock_response.headers = {"Content-Length": "100"}
|
|
mock_response.raw = MagicMock()
|
|
mock_response.raise_for_status = MagicMock()
|
|
mock_get.return_value = mock_response
|
|
|
|
mock_tqdm.wrapattr.return_value.__enter__ = MagicMock()
|
|
mock_tqdm.wrapattr.return_value.__exit__ = MagicMock()
|
|
|
|
result = download_assets(asset)
|
|
assert result == asset.filename
|
|
mock_logger.info.assert_called_with("Downloading %s assets", asset.filename)
|
|
mock_md5.assert_called_once_with(
|
|
asset.filename, "0f5a4b98abf3e3973faf9e9260a7d876"
|
|
)
|