# ------------------------------------------------------------------------ # RF-DETR # Copyright (c) 2025 Roboflow. All Rights Reserved. # Licensed under the Apache License, Version 2.0 [see LICENSE for details] # ------------------------------------------------------------------------ from types import ModuleType from unittest.mock import Mock, patch import pytest from rfdetr.assets import ModelWeightAsset, ModelWeights from rfdetr.assets.model_weights import download_pretrain_weights from rfdetr.platform import _IS_RFDETR_PLUS_AVAILABLE # Module-level fixture for common file operation mocks @pytest.fixture def mock_file_operations(): """Mock file operations to avoid actual file I/O.""" with ( patch("rfdetr.assets.model_weights.os.path.exists") as mock_exists, patch("rfdetr.assets.model_weights._download_file") as mock_download, patch("rfdetr.assets.model_weights._validate_file_md5") as mock_validate, ): # Default: file doesn't exist mock_exists.return_value = False # Default: MD5 validation passes mock_validate.return_value = True yield {"exists": mock_exists, "download": mock_download, "validate": mock_validate} class TestDownloadPretrainWeights: """Test download_pretrain_weights function with mocking for offline testing.""" def test_download_from_local_model_weights(self, mock_file_operations): """Test downloading a model from local ModelWeights.""" download_pretrain_weights("rf-detr-base.pth") # Should call download with correct URL and MD5 mock_file_operations["download"].assert_called_once() call_kwargs = mock_file_operations["download"].call_args[1] assert call_kwargs["filename"] == "rf-detr-base.pth" assert "rf-detr-base-coco.pth" in call_kwargs["url"] assert call_kwargs["expected_md5"] is not None # Should have MD5 hash assert len(call_kwargs["expected_md5"]) == 32 # Valid MD5 hash @pytest.mark.skipif(not _IS_RFDETR_PLUS_AVAILABLE, reason="rf-detr-plus not installed - skip priority test") def test_download_from_rfdetr_plus_when_available(self, mock_file_operations): """Test that rf-detr-plus URL is used for models absent from local ModelWeights. Verifies the priority contract: a model not registered in local ModelWeights but present in rf-detr-plus triggers _download_file with the plus URL, not a fallback or silent no-op. """ plus_url = "https://plus.example.com/plus-only-model.pth" plus_asset = ModelWeightAsset(filename="plus-only-model.pth", url=plus_url, md5_hash=None) with ( patch("rfdetr.assets.model_weights.ModelWeights.from_filename", return_value=None), patch("rfdetr_plus.assets.ModelWeights.from_filename", return_value=plus_asset), ): download_pretrain_weights("plus-only-model.pth") mock_file_operations["download"].assert_called_once() call_kwargs = mock_file_operations["download"].call_args[1] assert call_kwargs["url"] == plus_url def test_download_from_platform_models_fallback(self, mock_file_operations): """Test falling back to PLATFORM_MODELS when model not in ModelWeights.""" # Mock PLATFORM_MODELS mock_platform_module = Mock() mock_platform_module.PLATFORM_MODELS = {"legacy-model.pth": "https://legacy.com/model.pth"} with ( patch("rfdetr.assets.model_weights.ModelWeights.from_filename", return_value=None), patch.dict("sys.modules", {"rfdetr.platform": Mock(), "rfdetr.platform.downloads": mock_platform_module}), ): download_pretrain_weights("legacy-model.pth") # Should call download with legacy URL mock_file_operations["download"].assert_called_once() call_kwargs = mock_file_operations["download"].call_args[1] assert call_kwargs["url"] == "https://legacy.com/model.pth" assert call_kwargs["expected_md5"] is None # Platform models don't have MD5 def test_private_platform_models_registry_imports_without_eager_fallback(self): """Test that a private-only registry imports cleanly without evaluating PLATFORM_MODELS early.""" from importlib import reload import rfdetr.platform as platform_pkg import rfdetr.platform.downloads as platform_downloads mock_rfdetr_plus_pkg = ModuleType("rfdetr_plus") mock_rfdetr_plus_pkg.__path__ = [] mock_platforms_pkg = ModuleType("rfdetr_plus.models") mock_platforms_pkg.__path__ = [] mock_downloads_module = ModuleType("rfdetr_plus.models.downloads") mock_downloads_module._PLATFORM_MODELS = {"legacy-model.pth": "https://legacy.com/model.pth"} mock_platforms_pkg.downloads = mock_downloads_module mock_rfdetr_plus_pkg.models = mock_platforms_pkg with ( patch.dict( "sys.modules", { "rfdetr_plus": mock_rfdetr_plus_pkg, "rfdetr_plus.models": mock_platforms_pkg, "rfdetr_plus.models.downloads": mock_downloads_module, }, ), patch.object(platform_pkg, "_IS_RFDETR_PLUS_AVAILABLE", True), ): reloaded = reload(platform_downloads) assert reloaded.PLATFORM_MODELS == {"legacy-model.pth": "https://legacy.com/model.pth"} reload(platform_downloads) def test_file_exists_with_correct_md5(self, mock_file_operations): """Test that download is skipped if file exists with correct MD5.""" mock_file_operations["exists"].return_value = True mock_file_operations["validate"].return_value = True download_pretrain_weights("rf-detr-base.pth") # Should not download if file exists with correct hash mock_file_operations["download"].assert_not_called() def test_file_exists_with_incorrect_md5_warns_and_skips(self, mock_file_operations): """Test that file is NOT re-downloaded when MD5 is incorrect and redownload=False. This protects fine-tuned checkpoints that share the same filename as a registry model (e.g. rf-detr-nano.pth) from being silently overwritten. """ mock_file_operations["exists"].return_value = True mock_file_operations["validate"].return_value = False # Incorrect MD5 with patch("rfdetr.assets.model_weights.logger.warning") as mock_warning: download_pretrain_weights("rf-detr-base.pth") # Should NOT re-download — the user's file must be preserved mock_file_operations["download"].assert_not_called() mock_warning.assert_called_once() warning_msg = mock_warning.call_args[0][0] assert "incorrect MD5 hash" in warning_msg assert "skipping re-download to avoid overwriting it" in warning_msg def test_redownload_flag_forces_download(self, mock_file_operations): """Test that redownload=True forces re-download even if file exists.""" mock_file_operations["exists"].return_value = True mock_file_operations["validate"].return_value = True download_pretrain_weights("rf-detr-base.pth", redownload=True) # Should download despite file existing mock_file_operations["download"].assert_called_once() def test_redownload_flag_forces_download_despite_incorrect_md5(self, mock_file_operations): """Test that redownload=True triggers download even when MD5 is incorrect. Verifies the force-redownload path where the user explicitly wants to overwrite an existing file (e.g. a fine- tuned checkpoint) with the original registry weights. """ mock_file_operations["exists"].return_value = True mock_file_operations["validate"].return_value = False # Incorrect MD5 download_pretrain_weights("rf-detr-base.pth", redownload=True) # Should download because redownload=True overrides the skip-on-existing-file guard mock_file_operations["download"].assert_called_once() def test_validate_md5_disabled(self, mock_file_operations): """Test that MD5 validation can be disabled.""" download_pretrain_weights("rf-detr-base.pth", validate_md5=False) # Should pass expected_md5=None when validation is disabled call_kwargs = mock_file_operations["download"].call_args[1] assert call_kwargs["expected_md5"] is None @patch("rfdetr.assets.model_weights.ModelWeights.from_filename", return_value=None) def test_nonexistent_model_returns_early(self, mock_from_filename, mock_file_operations): """Test that function returns early for non-existent models.""" download_pretrain_weights("nonexistent-model.pth") # Should not attempt download mock_file_operations["download"].assert_not_called() def test_model_without_md5_hash(self, mock_file_operations): """Test downloading a model that has no MD5 hash.""" # Create a mock asset without MD5 mock_asset = ModelWeightAsset(filename="test-no-md5.pth", url="https://example.com/test.pth", md5_hash=None) with patch("rfdetr.assets.model_weights.ModelWeights.from_filename", return_value=mock_asset): download_pretrain_weights("test-no-md5.pth", validate_md5=True) # Should pass None for expected_md5 call_kwargs = mock_file_operations["download"].call_args[1] assert call_kwargs["expected_md5"] is None def test_file_exists_no_md5_skips_download(self, mock_file_operations): """Test that if file exists and no MD5 validation, download is skipped.""" mock_file_operations["exists"].return_value = True # Create a mock asset without MD5 mock_asset = ModelWeightAsset(filename="test-no-md5.pth", url="https://example.com/test.pth", md5_hash=None) with patch("rfdetr.assets.model_weights.ModelWeights.from_filename", return_value=mock_asset): download_pretrain_weights("test-no-md5.pth") # Should not download if file exists (no MD5 to validate) mock_file_operations["download"].assert_not_called() class TestDownloadIntegration: """Integration tests for the complete download flow.""" @pytest.mark.parametrize("model", [pytest.param(m, id=m.filename) for m in ModelWeights]) def test_all_models_have_valid_md5_format(self, model: ModelWeightAsset) -> None: """Test that MD5 hashes are valid when present (prevent typos).""" # MD5 should be None or valid 32-char hex string if model.md5_hash is not None: assert len(model.md5_hash) == 32, f"{model.filename} has invalid MD5 length: {len(model.md5_hash)}" assert all(c in "0123456789abcdef" for c in model.md5_hash.lower()), ( f"{model.filename} has invalid MD5 characters" ) def test_from_filename_bidirectional_lookup(self): """Test that from_filename correctly maps back to enum values.""" from rfdetr.assets.model_weights import ModelWeights # Pick a known model original = ModelWeights.RF_DETR_BASE # Look it up by filename asset = ModelWeights.from_filename(original.filename) # Should return the exact same asset assert asset is not None assert asset.filename == original.filename assert asset.url == original.url assert asset.md5_hash == original.md5_hash @patch("rfdetr.assets.model_weights.os.path.exists") @patch("rfdetr.assets.model_weights._validate_file_md5") @patch("rfdetr.assets.model_weights._download_file") def test_download_flow_for_real_model(self, mock_download, mock_validate, mock_exists): """Test the complete download flow for a real model.""" mock_exists.return_value = False mock_validate.return_value = True # Download a real model (mocked network) download_pretrain_weights("rf-detr-base.pth") # Verify download was called with correct parameters mock_download.assert_called_once() call_kwargs = mock_download.call_args[1] assert call_kwargs["filename"] == "rf-detr-base.pth" assert "storage.googleapis.com/rfdetr" in call_kwargs["url"] assert call_kwargs["expected_md5"] == "b4d3ce46099eaed50626ede388caf979" class TestDownloadErrorHandling: """Test error handling in download mechanism.""" @patch("rfdetr.assets.model_weights.os.path.exists") @patch("rfdetr.assets.model_weights._download_file") def test_handles_missing_rfdetr_plus_gracefully(self, mock_download, mock_exists): """Test that missing rf-detr-plus is handled gracefully.""" mock_exists.return_value = False # Should not raise an error if rf-detr-plus is not installed download_pretrain_weights("rf-detr-base.pth") # Should still download from local ModelWeights mock_download.assert_called_once() @patch("rfdetr.assets.model_weights.ModelWeights.from_filename", return_value=None) @patch("rfdetr.assets.model_weights._download_file") @patch("rfdetr.assets.model_weights.os.path.exists") def test_handles_missing_platform_models_gracefully(self, mock_exists, mock_download, mock_from_filename): """Test that missing platform models is handled gracefully.""" mock_exists.return_value = False # Should return early without raising error download_pretrain_weights("unknown-model.pth") # Should not attempt download mock_download.assert_not_called() @patch("rfdetr.assets.model_weights._validate_file_md5") @patch("rfdetr.assets.model_weights._download_file") @patch("rfdetr.assets.model_weights.os.path.exists") @patch("rfdetr.assets.model_weights.logger") def test_logs_info_messages(self, mock_logger, mock_exists, mock_download, mock_validate): """Test that appropriate log messages are generated.""" mock_exists.return_value = False mock_validate.return_value = True download_pretrain_weights("rf-detr-base.pth") # Should log download message mock_logger.info.assert_called() log_message = mock_logger.info.call_args[0][0] assert "rf-detr-base.pth" in log_message @patch("rfdetr.assets.model_weights._download_file") @patch("rfdetr.assets.model_weights._validate_file_md5") @patch("rfdetr.assets.model_weights.os.path.exists") @patch("rfdetr.assets.model_weights.logger") def test_logs_warning_on_incorrect_md5(self, mock_logger, mock_exists, mock_validate, mock_download): """Test that warning is logged when MD5 is incorrect and no re-download occurs.""" mock_exists.return_value = True mock_validate.return_value = False download_pretrain_weights("rf-detr-base.pth") # Should log warning about incorrect MD5 mock_logger.warning.assert_called() warning_message = mock_logger.warning.call_args[0][0] assert "incorrect MD5 hash" in warning_message # Must NOT re-download — fine-tuned checkpoints should be preserved mock_download.assert_not_called() @patch("rfdetr.assets.model_weights._download_file") @patch("rfdetr.assets.model_weights.os.path.exists") def test_absolute_path_resolves_to_known_model(self, mock_exists, mock_download): """Absolute paths like /content/rf-detr-base.pth must still match the registry. Regression test: previously ModelWeights.from_filename received the full path instead of the basename, so it returned None and the download was silently skipped. """ mock_exists.return_value = False download_pretrain_weights("/content/rf-detr-base.pth") # Must have attempted a download — not silently returned mock_download.assert_called_once() call_kwargs = mock_download.call_args[1] assert call_kwargs["filename"] == "/content/rf-detr-base.pth" assert "rf-detr-base-coco.pth" in call_kwargs["url"] @patch("rfdetr.assets.model_weights._download_file") @patch("rfdetr.assets.model_weights.os.path.exists") def test_nested_absolute_path_resolves_to_known_model(self, mock_exists, mock_download): """Nested paths like /workspace/models/rf-detr-base.pth also resolve.""" mock_exists.return_value = False download_pretrain_weights("/workspace/models/rf-detr-base.pth") mock_download.assert_called_once() call_kwargs = mock_download.call_args[1] assert call_kwargs["filename"] == "/workspace/models/rf-detr-base.pth" # --------------------------------------------------------------------------- # maybe_download_pretrain_weights — RF_HOME cache-dir path resolution # --------------------------------------------------------------------------- class TestMaybeDownloadPretrainWeightsCacheDir: """Verify that RFDETR.maybe_download_pretrain_weights resolves paths via RF_HOME.""" def _make_rfdetr(self, pretrain_weights): """Return an RFDETR shell backed by a fully validated RFDETRBaseConfig. Uses RFDETRBaseConfig (which supplies required field defaults) so the expand_path field validator on pretrain_weights is exercised end-to-end. Args: pretrain_weights: Raw value to pass to RFDETRBaseConfig; the pydantic validator resolves it before assigning to model_config.pretrain_weights. """ from rfdetr.config import RFDETRBaseConfig from rfdetr.detr import RFDETR model = object.__new__(RFDETR) model.model_config = RFDETRBaseConfig(pretrain_weights=pretrain_weights) return model def test_bare_filename_resolved_to_rf_home(self, monkeypatch, tmp_path): """Bare filename (no directory separator) is joined with RF_HOME before download.""" monkeypatch.setenv("RF_HOME", str(tmp_path)) downloaded = [] monkeypatch.setattr("rfdetr.detr.download_pretrain_weights", lambda p, **kw: downloaded.append(p)) self._make_rfdetr("rf-detr-base.pth").maybe_download_pretrain_weights() assert downloaded == [str(tmp_path / "rf-detr-base.pth")] def test_path_with_directory_used_as_is(self, monkeypatch, tmp_path): """Path containing a directory separator is not modified by RF_HOME.""" monkeypatch.setenv("RF_HOME", str(tmp_path / "should_not_be_used")) downloaded = [] monkeypatch.setattr("rfdetr.detr.download_pretrain_weights", lambda p, **kw: downloaded.append(p)) explicit = str(tmp_path / "custom" / "my_weights.pth") self._make_rfdetr(explicit).maybe_download_pretrain_weights() assert downloaded == [explicit] def test_none_pretrain_weights_skips_download(self, monkeypatch): """None pretrain_weights returns without calling download.""" called = [] monkeypatch.setattr("rfdetr.detr.download_pretrain_weights", lambda *a, **kw: called.append(True)) self._make_rfdetr(None).maybe_download_pretrain_weights() assert called == [] def test_cache_dir_created_when_absent(self, monkeypatch, tmp_path): """RF_HOME directory is created if it does not already exist.""" cache_dir = tmp_path / "new_cache" monkeypatch.setenv("RF_HOME", str(cache_dir)) monkeypatch.setattr("rfdetr.detr.download_pretrain_weights", lambda *a, **kw: None) self._make_rfdetr("rf-detr-base.pth").maybe_download_pretrain_weights() assert cache_dir.is_dir()