mirrors/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2025-09-13 04:55:53 +00:00
Code Issues Packages Projects Releases Wiki Activity

Use pydantic.

Browse Source
This commit is contained in:
Robin Huang
2024-08-08 13:06:26 -07:00
parent db1ce51fdf
commit 7461e8eb00
3 changed files with 57 additions and 34 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
tests-unit/prompt_server_test/download_models_test.py
View File

@@ -84,13 +84,19 @@ async def test_download_model_success():
# Check initial call
mock_progress_callback.assert_any_call(
'checkpoints/model.bin',
DownloadModelStatus(DownloadStatusType.PENDING, 0, "Starting download of model.bin", False)
DownloadModelStatus(status=DownloadStatusType.PENDING,
progress_percentage= 0,
message="Starting download of model.bin",
already_existed= False)
)
# Check final call
mock_progress_callback.assert_any_call(
'checkpoints/model.bin',
DownloadModelStatus(DownloadStatusType.COMPLETED, 100, "Successfully downloaded model.bin", False)
DownloadModelStatus(status=DownloadStatusType.COMPLETED,
progress_percentage=100,
message="Successfully downloaded model.bin",
already_existed= False)
)
# Verify file writing
@@ -204,7 +210,10 @@ async def test_check_file_exists_when_file_exists(tmp_path):
mock_callback.assert_called_once_with(
"test/existing_model.bin",
DownloadModelStatus(DownloadStatusType.COMPLETED, 100, "existing_model.bin already exists", already_existed=True)
DownloadModelStatus(status=DownloadStatusType.COMPLETED,
progress_percentage=100,
message="existing_model.bin already exists",
already_existed=True)
)
@pytest.mark.asyncio
@@ -237,7 +246,7 @@ async def test_track_download_progress_no_content_length():
# Check that progress was reported even without knowing the total size
mock_callback.assert_any_call(
'models/model.bin',
DownloadModelStatus(DownloadStatusType.IN_PROGRESS, 0, "Downloading model.bin", already_existed=False)
DownloadModelStatus(status=DownloadStatusType.IN_PROGRESS, progress_percentage= 0, message="Downloading model.bin", already_existed=False)
)
@pytest.mark.asyncio