diff --git a/test/assets/__init__.py b/test/assets/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/test/assets/test_downloader.py b/test/assets/test_downloader.py new file mode 100644 index 00000000..835826f1 --- /dev/null +++ b/test/assets/test_downloader.py @@ -0,0 +1,141 @@ +from unittest.mock import MagicMock, mock_open, patch + +import pytest + +from supervision.assets.downloader import download_assets, is_md5_hash_matching +from supervision.assets.list import VideoAssets + + +class TestMD5HashMatching: + def test_file_exists_matching_hash(self): + """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): + """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): + """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("builtins.print") + @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_print): + """Test download_assets when file already exists and is valid.""" + filename = "vehicles.mp4" + result = download_assets(filename) + assert result == filename + mock_print.assert_called_with(f"{filename} asset download complete. \n") + + @patch("supervision.assets.downloader.download_assets", return_value="vehicles.mp4") + @patch("os.remove") + @patch("supervision.assets.downloader.is_md5_hash_matching", return_value=False) + @patch("pathlib.Path.exists", return_value=True) + def test_already_exists_but_corrupted( + self, mock_exists, mock_md5, mock_remove, mock_recursive + ): + """Test download_assets when file exists but is corrupted (re-downloads).""" + filename = "vehicles.mp4" + result = download_assets(filename) + assert result == filename + mock_recursive.assert_called_with(filename) + + @patch("builtins.print") + @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_print, + ): + """Test download_assets downloading a new 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() + + result = download_assets(filename) + assert result == filename + mock_print.assert_called_with(f"Downloading {filename} assets \n") + mock_get.assert_called_once() + mock_response.raise_for_status.assert_called_once_with() + mock_copyfileobj.assert_called_once() + + @patch("pathlib.Path.exists", return_value=False) + def test_invalid_asset(self, mock_exists): + """Test download_assets with invalid asset name.""" + invalid_filename = "invalid.mp4" + + with pytest.raises(ValueError) as exc_info: + download_assets(invalid_filename) + + assert "Invalid asset" in str(exc_info.value) + assert "vehicles.mp4" in str(exc_info.value) + + @patch("builtins.print") + @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_enum( + self, + mock_exists, + mock_get, + mock_tqdm, + mock_copyfileobj, + mock_mkdir, + mock_open_file, + mock_print, + ): + """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.value diff --git a/test/assets/test_list.py b/test/assets/test_list.py new file mode 100644 index 00000000..bf5dc47c --- /dev/null +++ b/test/assets/test_list.py @@ -0,0 +1,48 @@ +from supervision.assets.list import BASE_VIDEO_URL, VIDEO_ASSETS, VideoAssets + + +def test_video_assets_list(): + """Test that VideoAssets.list() returns all video filenames.""" + expected_filenames = [ + "vehicles.mp4", + "milk-bottling-plant.mp4", + "vehicles-2.mp4", + "grocery-store.mp4", + "subway.mp4", + "market-square.mp4", + "people-walking.mp4", + "beach-1.mp4", + "basketball-1.mp4", + "skiing.mp4", + ] + assert VideoAssets.list() == expected_filenames + + +def test_video_assets_enum_values(): + """Test that VideoAssets enum members have correct values.""" + assert VideoAssets.VEHICLES.value == "vehicles.mp4" + assert VideoAssets.MILK_BOTTLING_PLANT.value == "milk-bottling-plant.mp4" + assert VideoAssets.VEHICLES_2.value == "vehicles-2.mp4" + assert VideoAssets.GROCERY_STORE.value == "grocery-store.mp4" + assert VideoAssets.SUBWAY.value == "subway.mp4" + assert VideoAssets.MARKET_SQUARE.value == "market-square.mp4" + assert VideoAssets.PEOPLE_WALKING.value == "people-walking.mp4" + assert VideoAssets.BEACH.value == "beach-1.mp4" + assert VideoAssets.BASKETBALL.value == "basketball-1.mp4" + assert VideoAssets.SKIING.value == "skiing.mp4" + + +def test_video_assets_dict_keys(): + """Test that VIDEO_ASSETS has all VideoAssets as keys.""" + expected_keys = {asset.value for asset in VideoAssets} + assert set(VIDEO_ASSETS.keys()) == expected_keys + + +def test_video_assets_dict_values(): + """Test that VIDEO_ASSETS values are tuples of (url, md5_hash).""" + for filename, (url, md5_hash) in VIDEO_ASSETS.items(): + assert isinstance(url, str) + assert url.startswith(BASE_VIDEO_URL) + assert url.endswith(filename) + assert isinstance(md5_hash, str) + assert len(md5_hash) == 32 # MD5 hash length