From 0df1e732207949c1e6bbf6f146059fc1d100fa40 Mon Sep 17 00:00:00 2001 From: "Alex \"mcmonkey\" Goodwin" Date: Sun, 22 Sep 2024 17:00:49 +0900 Subject: [PATCH] rewrite test for create_model_path --- .../prompt_server_test/download_models_test.py | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/tests-unit/prompt_server_test/download_models_test.py b/tests-unit/prompt_server_test/download_models_test.py index 8f633f8cf..bd41745fa 100644 --- a/tests-unit/prompt_server_test/download_models_test.py +++ b/tests-unit/prompt_server_test/download_models_test.py @@ -202,19 +202,21 @@ async def test_download_model_invalid_folder_path(): assert result.status == 'error' assert result.already_existed is False -# For create_model_path function def test_create_model_path(tmp_path, monkeypatch): - mock_models_dir = tmp_path / "models" - monkeypatch.setattr('folder_paths.models_dir', str(mock_models_dir)) + model_name = "model.safetensors" + folder_path = os.path.join(tmp_path, "mock_dir") - model_name = "test_model.sft" - model_directory = "test_dir" + file_path = create_model_path(model_name, folder_path) - file_path = create_model_path(model_name, model_directory, mock_models_dir) - - assert file_path == str(mock_models_dir / model_directory / model_name) + assert file_path == os.path.join(folder_path, "model.safetensors") assert os.path.exists(os.path.dirname(file_path)) + with pytest.raises(Exception, match="Invalid model directory"): + create_model_path("../path_traversal.safetensors", folder_path) + + with pytest.raises(Exception, match="Invalid model directory"): + create_model_path("/etc/some_root_path", folder_path) + @pytest.mark.asyncio async def test_check_file_exists_when_file_exists(tmp_path):