67 lines
2.4 KiB
Python
67 lines
2.4 KiB
Python
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()
|