feat(models): add bulk model deletion controls
This commit is contained in:
@@ -0,0 +1,66 @@
|
||||
from mapper import ModelMapper
|
||||
|
||||
|
||||
def _configure_model_paths(monkeypatch, tmp_path, filenames):
|
||||
import folder_paths
|
||||
|
||||
models_root = tmp_path / "models"
|
||||
models_root.mkdir()
|
||||
for filename in filenames:
|
||||
target = models_root / filename
|
||||
target.parent.mkdir(parents=True, exist_ok=True)
|
||||
target.write_bytes(b"model")
|
||||
|
||||
monkeypatch.setattr(folder_paths, "folder_names_and_paths", {"checkpoints": ([], {})})
|
||||
monkeypatch.setattr(folder_paths, "get_filename_list", lambda model_type: filenames if model_type == "checkpoints" else [])
|
||||
monkeypatch.setattr(folder_paths, "get_folder_paths", lambda model_type: [str(models_root)] if model_type == "checkpoints" else [])
|
||||
monkeypatch.setattr(folder_paths, "get_full_path", lambda model_type, filename: str(models_root / filename))
|
||||
return models_root
|
||||
|
||||
|
||||
def test_delete_models_removes_only_a_known_model(monkeypatch, tmp_path):
|
||||
models_root = _configure_model_paths(monkeypatch, tmp_path, ["nested/model.safetensors"])
|
||||
mapper = ModelMapper()
|
||||
|
||||
result = mapper.delete_models([{
|
||||
"model_type": "checkpoints",
|
||||
"model_name": "nested/model.safetensors",
|
||||
}])
|
||||
|
||||
assert result == [{
|
||||
"model_type": "checkpoints",
|
||||
"model_name": "nested/model.safetensors",
|
||||
"status": "deleted",
|
||||
}]
|
||||
assert not (models_root / "nested/model.safetensors").exists()
|
||||
|
||||
|
||||
def test_delete_models_rejects_unknown_name(monkeypatch, tmp_path):
|
||||
models_root = _configure_model_paths(monkeypatch, tmp_path, ["known.safetensors"])
|
||||
mapper = ModelMapper()
|
||||
|
||||
result = mapper.delete_models([{
|
||||
"model_type": "checkpoints",
|
||||
"model_name": "../outside.safetensors",
|
||||
}])
|
||||
|
||||
assert result[0]["status"] == "error"
|
||||
assert (models_root / "known.safetensors").exists()
|
||||
|
||||
|
||||
def test_delete_models_rejects_path_outside_model_root(monkeypatch, tmp_path):
|
||||
models_root = _configure_model_paths(monkeypatch, tmp_path, ["known.safetensors"])
|
||||
outside = tmp_path / "outside.safetensors"
|
||||
outside.write_bytes(b"do not remove")
|
||||
import folder_paths
|
||||
monkeypatch.setattr(folder_paths, "get_full_path", lambda *_: str(outside))
|
||||
mapper = ModelMapper()
|
||||
|
||||
result = mapper.delete_models([{
|
||||
"model_type": "checkpoints",
|
||||
"model_name": "known.safetensors",
|
||||
}])
|
||||
|
||||
assert result[0]["status"] == "error"
|
||||
assert (models_root / "known.safetensors").exists()
|
||||
assert outside.exists()
|
||||
Reference in New Issue
Block a user