alphard/ComfyUI-Lora-Manager
1
0
Fork 0
mirror of https://github.com/willmiao/ComfyUI-Lora-Manager.git synced 2026-03-21 21:22:11 -03:00
Code Issues Packages Projects Releases Wiki Activity

feat(example-images): add stop control for download panel

Browse Source
This commit is contained in:
pixelpaws
2025-10-14 22:18:57 +08:00
parent a3070f8d82
commit 52e671638b
18 changed files with 389 additions and 57 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
locales/de.json
View File

@@ -1245,6 +1245,8 @@
"pauseFailed": "Fehler beim Pausieren des Downloads: {error}",
"downloadResumed": "Download fortgesetzt",
"resumeFailed": "Fehler beim Fortsetzen des Downloads: {error}",
"downloadStopped": "Download abgebrochen",
"stopFailed": "Download konnte nicht abgebrochen werden: {error}",
"deleted": "Beispielbild gelöscht",
"deleteFailed": "Fehler beim Löschen des Beispielbilds",
"setPreviewFailed": "Fehler beim Setzen des Vorschaubilds"

2
locales/en.json
View File

@@ -1245,6 +1245,8 @@
"pauseFailed": "Failed to pause download: {error}",
"downloadResumed": "Download resumed",
"resumeFailed": "Failed to resume download: {error}",
"downloadStopped": "Download cancelled",
"stopFailed": "Failed to cancel download: {error}",
"deleted": "Example image deleted",
"deleteFailed": "Failed to delete example image",
"setPreviewFailed": "Failed to set preview image"

2
locales/es.json
View File

@@ -1245,6 +1245,8 @@
"pauseFailed": "Error al pausar descarga: {error}",
"downloadResumed": "Descarga reanudada",
"resumeFailed": "Error al reanudar descarga: {error}",
"downloadStopped": "Descarga cancelada",
"stopFailed": "Error al cancelar descarga: {error}",
"deleted": "Imagen de ejemplo eliminada",
"deleteFailed": "Error al eliminar imagen de ejemplo",
"setPreviewFailed": "Error al establecer imagen de vista previa"

2
locales/fr.json
View File

@@ -1245,6 +1245,8 @@
"pauseFailed": "Échec de la mise en pause du téléchargement : {error}",
"downloadResumed": "Téléchargement repris",
"resumeFailed": "Échec de la reprise du téléchargement : {error}",
"downloadStopped": "Téléchargement annulé",
"stopFailed": "Échec de l'annulation du téléchargement : {error}",
"deleted": "Image d'exemple supprimée",
"deleteFailed": "Échec de la suppression de l'image d'exemple",
"setPreviewFailed": "Échec de la définition de l'image d'aperçu"

2
locales/he.json
View File

@@ -1245,6 +1245,8 @@
"pauseFailed": "השהיית ההורדה נכשלה: {error}",
"downloadResumed": "ההורדה חודשה",
"resumeFailed": "חידוש ההורדה נכשל: {error}",
"downloadStopped": "ההורדה בוטלה",
"stopFailed": "נכשל בביטול ההורדה: {error}",
"deleted": "תמונת הדוגמה נמחקה",
"deleteFailed": "מחיקת תמונת הדוגמה נכשלה",
"setPreviewFailed": "הגדרת תמונת התצוגה המקדימה נכשלה"

2
locales/ja.json
View File

@@ -1245,6 +1245,8 @@
"pauseFailed": "ダウンロードの一時停止に失敗しました：{error}",
"downloadResumed": "ダウンロードが再開されました",
"resumeFailed": "ダウンロードの再開に失敗しました：{error}",
"downloadStopped": "ダウンロードをキャンセルしました",
"stopFailed": "ダウンロードのキャンセルに失敗しました：{error}",
"deleted": "例画像が削除されました",
"deleteFailed": "例画像の削除に失敗しました",
"setPreviewFailed": "プレビュー画像の設定に失敗しました"

2
locales/ko.json
View File

@@ -1245,6 +1245,8 @@
"pauseFailed": "다운로드 일시정지 실패: {error}",
"downloadResumed": "다운로드가 재개되었습니다",
"resumeFailed": "다운로드 재개 실패: {error}",
"downloadStopped": "다운로드가 취소되었습니다",
"stopFailed": "다운로드 취소 실패: {error}",
"deleted": "예시 이미지가 삭제되었습니다",
"deleteFailed": "예시 이미지 삭제 실패",
"setPreviewFailed": "미리보기 이미지 설정 실패"

2
locales/ru.json
View File

@@ -1245,6 +1245,8 @@
"pauseFailed": "Не удалось приостановить загрузку: {error}",
"downloadResumed": "Загрузка возобновлена",
"resumeFailed": "Не удалось возобновить загрузку: {error}",
"downloadStopped": "Загрузка отменена",
"stopFailed": "Не удалось отменить загрузку: {error}",
"deleted": "Пример изображения удален",
"deleteFailed": "Не удалось удалить пример изображения",
"setPreviewFailed": "Не удалось установить превью изображение"

2
locales/zh-CN.json
View File

@@ -1245,6 +1245,8 @@
"pauseFailed": "暂停下载失败：{error}",
"downloadResumed": "下载已恢复",
"resumeFailed": "恢复下载失败：{error}",
"downloadStopped": "下载已取消",
"stopFailed": "取消下载失败：{error}",
"deleted": "示例图片已删除",
"deleteFailed": "删除示例图片失败",
"setPreviewFailed": "设置预览图片失败"

2
locales/zh-TW.json
View File

@@ -1245,6 +1245,8 @@
"pauseFailed": "暫停下載失敗：{error}",
"downloadResumed": "下載已恢復",
"resumeFailed": "恢復下載失敗：{error}",
"downloadStopped": "下載已取消",
"stopFailed": "取消下載失敗：{error}",
"deleted": "範例圖片已刪除",
"deleteFailed": "刪除範例圖片失敗",
"setPreviewFailed": "設定預覽圖片失敗"

1
py/routes/example_images_route_registrar.py
View File

@@ -22,6 +22,7 @@ ROUTE_DEFINITIONS: tuple[RouteDefinition, ...] = (
RouteDefinition("GET", "/api/lm/example-images-status", "get_example_images_status"),
RouteDefinition("POST", "/api/lm/pause-example-images", "pause_example_images"),
RouteDefinition("POST", "/api/lm/resume-example-images", "resume_example_images"),
RouteDefinition("POST", "/api/lm/stop-example-images", "stop_example_images"),
RouteDefinition("POST", "/api/lm/open-example-images-folder", "open_example_images_folder"),
RouteDefinition("GET", "/api/lm/example-image-files", "get_example_image_files"),
RouteDefinition("GET", "/api/lm/has-example-images", "has_example_images"),

8
py/routes/handlers/example_images_handlers.py
View File

@@ -68,6 +68,13 @@ class ExampleImagesDownloadHandler:
except DownloadNotRunningError as exc:
return web.json_response({'success': False, 'error': str(exc)}, status=400)
async def stop_example_images(self, request: web.Request) -> web.StreamResponse:
try:
result = await self._download_manager.stop_download(request)
return web.json_response(result)
except DownloadNotRunningError as exc:
return web.json_response({'success': False, 'error': str(exc)}, status=400)
async def force_download_example_images(self, request: web.Request) -> web.StreamResponse:
try:
payload = await request.json()
@@ -149,6 +156,7 @@ class ExampleImagesHandlerSet:
"get_example_images_status": self.download.get_example_images_status,
"pause_example_images": self.download.pause_example_images,
"resume_example_images": self.download.resume_example_images,
"stop_example_images": self.download.stop_example_images,
"force_download_example_images": self.download.force_download_example_images,
"import_example_images": self.management.import_example_images,
"delete_example_image": self.management.delete_example_image,

160
py/utils/example_images_download_manager.py
View File

@@ -105,6 +105,7 @@ class DownloadManager:
self._progress = _DownloadProgress()
self._ws_manager = ws_manager
self._state_lock = state_lock or asyncio.Lock()
self._stop_requested = False
def _resolve_output_dir(self, library_name: str | None = None) -> str:
base_path = get_settings_manager().get('example_images_path')
@@ -145,6 +146,7 @@ class DownloadManager:
raise DownloadConfigurationError('Example images path not configured in settings')
self._progress.reset()
self._stop_requested = False
self._progress['status'] = 'running'
self._progress['start_time'] = time.time()
self._progress['end_time'] = None
@@ -268,6 +270,27 @@ class DownloadManager:
'message': 'Download resumed'
}
async def stop_download(self, request):
"""Stop the example images download after the current model completes."""
async with self._state_lock:
if not self._is_downloading:
raise DownloadNotRunningError()
if self._progress['status'] in {'completed', 'error', 'stopped'}:
raise DownloadNotRunningError()
if self._progress['status'] != 'stopping':
self._stop_requested = True
self._progress['status'] = 'stopping'
await self._broadcast_progress(status='stopping')
return {
'success': True,
'message': 'Download stopping'
}
async def _download_all_example_images(
self,
output_dir,
@@ -311,6 +334,12 @@ class DownloadManager:
# Process each model
for i, (scanner_type, model, scanner) in enumerate(all_models):
async with self._state_lock:
current_status = self._progress['status']
if current_status not in {'running', 'paused', 'stopping'}:
break
# Main logic for processing model is here, but actual operations are delegated to other classes
was_remote_download = await self._process_model(
scanner_type,
@@ -324,21 +353,56 @@ class DownloadManager:
# Update progress
self._progress['completed'] += 1
await self._broadcast_progress(status='running')
async with self._state_lock:
current_status = self._progress['status']
should_stop = self._stop_requested and current_status == 'stopping'
broadcast_status = 'running' if current_status == 'running' else current_status
await self._broadcast_progress(status=broadcast_status)
if should_stop:
break
# Only add delay after remote download of models, and not after processing the last model
if was_remote_download and i < len(all_models) - 1 and self._progress['status'] == 'running':
if (
was_remote_download
and i < len(all_models) - 1
and current_status == 'running'
):
await asyncio.sleep(delay)
# Mark as completed
self._progress['status'] = 'completed'
self._progress['end_time'] = time.time()
logger.debug(
"Example images download completed: %s/%s models processed",
self._progress['completed'],
self._progress['total'],
)
await self._broadcast_progress(status='completed')
async with self._state_lock:
if self._stop_requested and self._progress['status'] == 'stopping':
self._progress['status'] = 'stopped'
self._progress['end_time'] = time.time()
self._stop_requested = False
final_status = 'stopped'
elif self._progress['status'] not in {'error', 'stopped'}:
self._progress['status'] = 'completed'
self._progress['end_time'] = time.time()
self._stop_requested = False
final_status = 'completed'
else:
final_status = self._progress['status']
self._stop_requested = False
if self._progress['end_time'] is None:
self._progress['end_time'] = time.time()
if final_status == 'completed':
logger.debug(
"Example images download completed: %s/%s models processed",
self._progress['completed'],
self._progress['total'],
)
elif final_status == 'stopped':
logger.debug(
"Example images download stopped: %s/%s models processed",
self._progress['completed'],
self._progress['total'],
)
await self._broadcast_progress(status=final_status)
except Exception as e:
error_msg = f"Error during example images download: {str(e)}"
@@ -360,6 +424,7 @@ class DownloadManager:
async with self._state_lock:
self._is_downloading = False
self._download_task = None
self._stop_requested = False
async def _process_model(
self,
@@ -378,7 +443,7 @@ class DownloadManager:
await asyncio.sleep(1)
# Check if download should continue
if self._progress['status'] != 'running':
if self._progress['status'] not in {'running', 'stopping'}:
logger.info(f"Download stopped: {self._progress['status']}")
return False # Return False to indicate no remote download happened
@@ -567,6 +632,7 @@ class DownloadManager:
raise DownloadConfigurationError('Example images path not configured in settings')
self._progress.reset()
self._stop_requested = False
self._progress['total'] = len(model_hashes)
self._progress['status'] = 'running'
self._progress['start_time'] = time.time()
@@ -588,10 +654,15 @@ class DownloadManager:
async with self._state_lock:
self._is_downloading = False
final_status = self._progress['status']
message = 'Force download completed'
if final_status == 'stopped':
message = 'Force download stopped'
return {
'success': True,
'message': 'Force download completed',
'message': message,
'result': result
}
@@ -649,6 +720,12 @@ class DownloadManager:
# Process each model
success_count = 0
for i, (scanner_type, model, scanner) in enumerate(models_to_process):
async with self._state_lock:
current_status = self._progress['status']
if current_status not in {'running', 'paused', 'stopping'}:
break
# Force process this model regardless of previous status
was_successful = await self._process_specific_model(
scanner_type,
@@ -666,24 +743,57 @@ class DownloadManager:
# Update progress
self._progress['completed'] += 1
async with self._state_lock:
current_status = self._progress['status']
should_stop = self._stop_requested and current_status == 'stopping'
broadcast_status = 'running' if current_status == 'running' else current_status
# Send progress update via WebSocket
await self._broadcast_progress(status='running')
await self._broadcast_progress(status=broadcast_status)
if should_stop:
break
# Only add delay after remote download, and not after processing the last model
if was_successful and i < len(models_to_process) - 1 and self._progress['status'] == 'running':
if (
was_successful
and i < len(models_to_process) - 1
and current_status == 'running'
):
await asyncio.sleep(delay)
# Mark as completed
self._progress['status'] = 'completed'
self._progress['end_time'] = time.time()
logger.debug(
"Forced example images download completed: %s/%s models processed",
self._progress['completed'],
self._progress['total'],
)
async with self._state_lock:
if self._stop_requested and self._progress['status'] == 'stopping':
self._progress['status'] = 'stopped'
self._progress['end_time'] = time.time()
self._stop_requested = False
final_status = 'stopped'
elif self._progress['status'] not in {'error', 'stopped'}:
self._progress['status'] = 'completed'
self._progress['end_time'] = time.time()
self._stop_requested = False
final_status = 'completed'
else:
final_status = self._progress['status']
self._stop_requested = False
if self._progress['end_time'] is None:
self._progress['end_time'] = time.time()
if final_status == 'completed':
logger.debug(
"Forced example images download completed: %s/%s models processed",
self._progress['completed'],
self._progress['total'],
)
elif final_status == 'stopped':
logger.debug(
"Forced example images download stopped: %s/%s models processed",
self._progress['completed'],
self._progress['total'],
)
# Send final progress via WebSocket
await self._broadcast_progress(status='completed')
await self._broadcast_progress(status=final_status)
return {
'total': self._progress['total'],
@@ -726,7 +836,7 @@ class DownloadManager:
await asyncio.sleep(1)
# Check if download should continue
if self._progress['status'] != 'running':
if self._progress['status'] not in {'running', 'stopping'}:
logger.info(f"Download stopped: {self._progress['status']}")
return False

116
static/js/managers/ExampleImagesManager.js
View File

@@ -13,8 +13,10 @@ export class ExampleImagesManager {
this.progressPanel = null;
this.isProgressPanelCollapsed = false;
this.pauseButton = null; // Store reference to the pause button
this.stopButton = null;
this.isMigrating = false; // Track migration state separately from downloading
this.hasShownCompletionToast = false; // Flag to track if completion toast has been shown
this.isStopping = false;
// Auto download properties
this.autoDownloadInterval = null;
@@ -52,12 +54,17 @@ export class ExampleImagesManager {
// Initialize progress panel button handlers
this.pauseButton = document.getElementById('pauseExampleDownloadBtn');
this.stopButton = document.getElementById('stopExampleDownloadBtn');
const collapseBtn = document.getElementById('collapseProgressBtn');
if (this.pauseButton) {
this.pauseButton.onclick = () => this.pauseDownload();
}
if (this.stopButton) {
this.stopButton.onclick = () => this.stopDownload();
}
if (collapseBtn) {
collapseBtn.onclick = () => this.toggleProgressPanel();
}
@@ -210,10 +217,14 @@ export class ExampleImagesManager {
updateDownloadButtonText() {
const btnTextElement = document.getElementById('exampleDownloadBtnText');
if (btnTextElement) {
if (this.isDownloading && this.isPaused) {
if (this.isStopping) {
btnTextElement.textContent = "Stopping...";
} else if (this.isDownloading && this.isPaused) {
btnTextElement.textContent = "Resume";
} else if (!this.isDownloading) {
btnTextElement.textContent = "Download";
} else {
btnTextElement.textContent = "Download";
}
}
}
@@ -243,12 +254,16 @@ export class ExampleImagesManager {
if (data.success) {
this.isDownloading = true;
this.isPaused = false;
this.isStopping = false;
this.hasShownCompletionToast = false; // Reset toast flag when starting new download
this.startTime = new Date();
this.updateUI(data.status);
this.showProgressPanel();
this.startProgressUpdates();
this.updateDownloadButtonText();
if (this.stopButton) {
this.stopButton.disabled = false;
}
showToast('toast.exampleImages.downloadStarted', {}, 'success');
// Close settings modal
@@ -263,7 +278,7 @@ export class ExampleImagesManager {
}
async pauseDownload() {
if (!this.isDownloading || this.isPaused) {
if (!this.isDownloading || this.isPaused || this.isStopping) {
return;
}
@@ -299,7 +314,7 @@ export class ExampleImagesManager {
}
async resumeDownload() {
if (!this.isDownloading || !this.isPaused) {
if (!this.isDownloading || !this.isPaused || this.isStopping) {
return;
}
@@ -334,6 +349,60 @@ export class ExampleImagesManager {
}
}
async stopDownload() {
if (this.isStopping) {
return;
}
if (!this.isDownloading) {
this.hideProgressPanel();
return;
}
this.isStopping = true;
this.isPaused = false;
this.updateDownloadButtonText();
if (this.stopButton) {
this.stopButton.disabled = true;
}
try {
const response = await fetch('/api/lm/stop-example-images', {
method: 'POST'
});
let data;
try {
data = await response.json();
} catch (parseError) {
data = { success: false, error: 'Invalid server response' };
}
if (response.ok && data.success) {
showToast('toast.exampleImages.downloadStopped', {}, 'info');
this.hideProgressPanel();
} else {
this.isStopping = false;
if (this.stopButton) {
this.stopButton.disabled = false;
}
const errorMessage = data && data.error ? data.error : 'Unknown error';
showToast('toast.exampleImages.stopFailed', { error: errorMessage }, 'error');
}
} catch (error) {
console.error('Failed to stop download:', error);
this.isStopping = false;
if (this.stopButton) {
this.stopButton.disabled = false;
}
const errorMessage = error && error.message ? error.message : 'Unknown error';
showToast('toast.exampleImages.stopFailed', { error: errorMessage }, 'error');
} finally {
this.updateDownloadButtonText();
}
}
startProgressUpdates() {
// Clear any existing interval
if (this.progressUpdateInterval) {
@@ -352,10 +421,22 @@ export class ExampleImagesManager {
const data = await response.json();
if (data.success) {
const currentStatus = data.status.status;
this.isDownloading = data.is_downloading;
this.isPaused = data.status.status === 'paused';
this.isPaused = currentStatus === 'paused';
this.isMigrating = data.is_migrating || false;
if (currentStatus === 'stopping') {
this.isStopping = true;
} else if (
!data.is_downloading ||
currentStatus === 'stopped' ||
currentStatus === 'completed' ||
currentStatus === 'error'
) {
this.isStopping = false;
}
// Update download button text
this.updateDownloadButtonText();
@@ -365,8 +446,11 @@ export class ExampleImagesManager {
// Download completed or failed
clearInterval(this.progressUpdateInterval);
this.progressUpdateInterval = null;
if (this.stopButton) {
this.stopButton.disabled = true;
}
if (data.status.status === 'completed' && !this.hasShownCompletionToast) {
if (currentStatus === 'completed' && !this.hasShownCompletionToast) {
const actionType = this.isMigrating ? 'migration' : 'download';
showToast('toast.downloads.imagesCompleted', { action: actionType }, 'success');
// Mark as shown to prevent duplicate toasts
@@ -375,10 +459,13 @@ export class ExampleImagesManager {
this.isMigrating = false;
// Hide the panel after a delay
setTimeout(() => this.hideProgressPanel(), 5000);
} else if (data.status.status === 'error') {
} else if (currentStatus === 'error') {
const actionType = this.isMigrating ? 'migration' : 'download';
showToast('toast.downloads.imagesFailed', { action: actionType }, 'error');
this.isMigrating = false;
} else if (currentStatus === 'stopped') {
this.hideProgressPanel();
this.isMigrating = false;
}
}
}
@@ -435,6 +522,10 @@ export class ExampleImagesManager {
this.pauseButton = document.getElementById('pauseExampleDownloadBtn');
}
if (!this.stopButton) {
this.stopButton = document.getElementById('stopExampleDownloadBtn');
}
if (this.pauseButton) {
// Check if the button already has the SVG elements
let hasProgressElements = !!this.pauseButton.querySelector('.mini-progress-circle');
@@ -462,6 +553,8 @@ export class ExampleImagesManager {
? () => this.resumeDownload()
: () => this.pauseDownload();
this.pauseButton.disabled = ['completed', 'error', 'stopped'].includes(status.status) || status.status === 'stopping';
// Update progress immediately
const progressBar = document.getElementById('downloadProgressBar');
if (progressBar) {
@@ -470,6 +563,15 @@ export class ExampleImagesManager {
}
}
if (this.stopButton) {
if (status.status === 'stopping' || this.isStopping) {
this.stopButton.disabled = true;
} else {
const canStop = ['running', 'paused'].includes(status.status);
this.stopButton.disabled = !canStop;
}
}
// Update title text
const titleElement = document.querySelector('.progress-panel-title');
if (titleElement) {
@@ -584,6 +686,8 @@ export class ExampleImagesManager {
case 'paused': return 'Paused';
case 'completed': return 'Completed';
case 'error': return 'Error';
case 'stopping': return 'Stopping';
case 'stopped': return 'Stopped';
default: return 'Initializing';
}
}

3
templates/components/progress_panel.html
View File

@@ -13,6 +13,9 @@
</svg>
<span class="progress-percent"></span>
</button>
<button id="stopExampleDownloadBtn" class="icon-button">
<i class="fas fa-stop"></i>
</button>
<button id="collapseProgressBtn" class="icon-button">
<i class="fas fa-chevron-down"></i>
</button>

13
tests/routes/test_example_images_route_registrar_handlers.py
View File

@@ -41,9 +41,11 @@ class StubDownloadManager:
def __init__(self) -> None:
self.pause_calls = 0
self.resume_calls = 0
self.stop_calls = 0
self.force_payloads: list[dict[str, Any]] = []
self.pause_error: Exception | None = None
self.resume_error: Exception | None = None
self.stop_error: Exception | None = None
self.force_error: Exception | None = None
async def get_status(self, request: web.Request) -> dict[str, Any]:
@@ -61,6 +63,12 @@ class StubDownloadManager:
raise self.resume_error
return {"success": True, "message": "resumed"}
async def stop_download(self, request: web.Request) -> dict[str, Any]:
self.stop_calls += 1
if self.stop_error:
raise self.stop_error
return {"success": True, "message": "stopping"}
async def start_force_download(self, payload: dict[str, Any]) -> dict[str, Any]:
self.force_payloads.append(payload)
if self.force_error:
@@ -193,17 +201,22 @@ async def test_pause_and_resume_return_client_errors_when_not_running():
async with registrar_app() as harness:
harness.download_manager.pause_error = DownloadNotRunningError()
harness.download_manager.resume_error = DownloadNotRunningError("Stopped")
harness.download_manager.stop_error = DownloadNotRunningError("Not running")
pause_response = await harness.client.post("/api/lm/pause-example-images")
resume_response = await harness.client.post("/api/lm/resume-example-images")
stop_response = await harness.client.post("/api/lm/stop-example-images")
assert pause_response.status == 400
assert resume_response.status == 400
assert stop_response.status == 400
pause_body = await _json(pause_response)
resume_body = await _json(resume_response)
stop_body = await _json(stop_response)
assert pause_body == {"success": False, "error": "No download in progress"}
assert resume_body == {"success": False, "error": "Stopped"}
assert stop_body == {"success": False, "error": "Not running"}
async def test_import_route_returns_validation_errors():

20
tests/routes/test_example_images_routes.py
View File

@@ -51,6 +51,10 @@ class StubDownloadManager:
self.calls.append(("resume_download", None))
return {"operation": "resume_download"}
async def stop_download(self, request: web.Request) -> dict:
self.calls.append(("stop_download", None))
return {"operation": "stop_download"}
async def start_force_download(self, payload: Any) -> dict:
self.calls.append(("start_force_download", payload))
return {"operation": "start_force_download", "payload": payload}
@@ -195,19 +199,23 @@ async def test_status_route_returns_manager_payload():
assert harness.download_manager.calls == [("get_status", {"detail": "true"})]
async def test_pause_and_resume_routes_delegate():
async def test_pause_resume_and_stop_routes_delegate():
async with example_images_app() as harness:
pause_response = await harness.client.post("/api/lm/pause-example-images")
resume_response = await harness.client.post("/api/lm/resume-example-images")
stop_response = await harness.client.post("/api/lm/stop-example-images")
assert pause_response.status == 200
assert await pause_response.json() == {"operation": "pause_download"}
assert resume_response.status == 200
assert await resume_response.json() == {"operation": "resume_download"}
assert stop_response.status == 200
assert await stop_response.json() == {"operation": "stop_download"}
assert harness.download_manager.calls[-2:] == [
assert harness.download_manager.calls[-3:] == [
("pause_download", None),
("resume_download", None),
("stop_download", None),
]
@@ -309,6 +317,10 @@ async def test_download_handler_methods_delegate() -> None:
self.calls.append(("resume_download", request))
return {"status": "running"}
async def stop_download(self, request) -> dict:
self.calls.append(("stop_download", request))
return {"status": "stopping"}
async def start_force_download(self, payload) -> dict:
self.calls.append(("start_force_download", payload))
return {"status": "force", "payload": payload}
@@ -342,6 +354,8 @@ async def test_download_handler_methods_delegate() -> None:
assert json.loads(pause_response.text) == {"status": "paused"}
resume_response = await handler.resume_example_images(request)
assert json.loads(resume_response.text) == {"status": "running"}
stop_response = await handler.stop_example_images(request)
assert json.loads(stop_response.text) == {"status": "stopping"}
force_response = await handler.force_download_example_images(request)
assert json.loads(force_response.text) == {"status": "force", "payload": {"foo": "bar"}}
@@ -350,6 +364,7 @@ async def test_download_handler_methods_delegate() -> None:
("get_status", request),
("pause_download", request),
("resume_download", request),
("stop_download", request),
("start_force_download", {"foo": "bar"}),
]
@@ -460,6 +475,7 @@ def test_handler_set_route_mapping_includes_all_handlers() -> None:
"get_example_images_status",
"pause_example_images",
"resume_example_images",
"stop_example_images",
"force_download_example_images",
"import_example_images",
"delete_example_image",

57
tests/utils/test_example_images_download_manager_unit.py
View File

@@ -1,6 +1,7 @@
from __future__ import annotations
import asyncio
import time
from typing import Any, Dict
import pytest
@@ -128,6 +129,59 @@ async def test_pause_and_resume_flow(monkeypatch: pytest.MonkeyPatch, tmp_path)
await asyncio.wait_for(task, timeout=1)
async def test_stop_download_transitions_to_stopped(monkeypatch: pytest.MonkeyPatch, tmp_path) -> None:
settings_manager = get_settings_manager()
settings_manager.settings["example_images_path"] = str(tmp_path)
settings_manager.settings["libraries"] = {"default": {}}
settings_manager.settings["active_library"] = "default"
ws_manager = RecordingWebSocketManager()
manager = download_module.DownloadManager(ws_manager=ws_manager)
started = asyncio.Event()
release = asyncio.Event()
async def fake_download(self, *_args):
started.set()
await release.wait()
async with self._state_lock:
if self._stop_requested and self._progress['status'] == 'stopping':
self._progress['status'] = 'stopped'
else:
self._progress['status'] = 'completed'
self._progress['end_time'] = time.time()
self._stop_requested = False
await self._broadcast_progress(status=self._progress['status'])
async with self._state_lock:
self._is_downloading = False
self._download_task = None
monkeypatch.setattr(
download_module.DownloadManager,
"_download_all_example_images",
fake_download,
)
await manager.start_download({})
await asyncio.wait_for(started.wait(), timeout=1)
stop_response = await manager.stop_download(object())
assert stop_response == {"success": True, "message": "Download stopping"}
assert manager._progress["status"] == "stopping"
task = manager._download_task
assert task is not None
release.set()
await asyncio.wait_for(task, timeout=1)
assert manager._progress["status"] == "stopped"
assert manager._is_downloading is False
assert manager._stop_requested is False
statuses = [payload["status"] for payload in ws_manager.payloads]
assert "stopping" in statuses
assert "stopped" in statuses
async def test_pause_or_resume_without_running_download(monkeypatch: pytest.MonkeyPatch) -> None:
manager = download_module.DownloadManager(ws_manager=RecordingWebSocketManager())
@@ -136,3 +190,6 @@ async def test_pause_or_resume_without_running_download(monkeypatch: pytest.Monk
with pytest.raises(download_module.DownloadNotRunningError):
await manager.resume_download(object())
with pytest.raises(download_module.DownloadNotRunningError):
await manager.stop_download(object())