mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-08-09 23:40:15 -03:00
feat(recipes): add recipe rematch WebSocket progress channel
This commit is contained in:
@@ -22,6 +22,8 @@ class WebSocketManager:
|
|||||||
self._auto_organize_progress: Optional[Dict[str, Any]] = None
|
self._auto_organize_progress: Optional[Dict[str, Any]] = None
|
||||||
# Add recipe repair progress tracking
|
# Add recipe repair progress tracking
|
||||||
self._recipe_repair_progress: Optional[Dict[str, Any]] = None
|
self._recipe_repair_progress: Optional[Dict[str, Any]] = None
|
||||||
|
# Add recipe rematch progress tracking
|
||||||
|
self._recipe_rematch_progress: Optional[Dict[str, Any]] = None
|
||||||
self._auto_organize_lock = asyncio.Lock()
|
self._auto_organize_lock = asyncio.Lock()
|
||||||
|
|
||||||
async def handle_connection(self, request: web.Request) -> web.WebSocketResponse:
|
async def handle_connection(self, request: web.Request) -> web.WebSocketResponse:
|
||||||
@@ -223,6 +225,30 @@ class WebSocketManager:
|
|||||||
status = self._recipe_repair_progress.get('status')
|
status = self._recipe_repair_progress.get('status')
|
||||||
return status in ['started', 'processing']
|
return status in ['started', 'processing']
|
||||||
|
|
||||||
|
async def broadcast_recipe_rematch_progress(self, data: Dict[str, Any]):
|
||||||
|
"""Broadcast recipe rematch progress to connected clients"""
|
||||||
|
# Store progress data in memory
|
||||||
|
self._recipe_rematch_progress = data
|
||||||
|
|
||||||
|
# Broadcast via WebSocket
|
||||||
|
await self.broadcast(data)
|
||||||
|
|
||||||
|
def get_recipe_rematch_progress(self) -> Optional[Dict[str, Any]]:
|
||||||
|
"""Get current recipe rematch progress"""
|
||||||
|
return self._recipe_rematch_progress
|
||||||
|
|
||||||
|
def cleanup_recipe_rematch_progress(self):
|
||||||
|
"""Clear recipe rematch progress data if it is in a finished state"""
|
||||||
|
if self._recipe_rematch_progress and self._recipe_rematch_progress.get('status') in ['completed', 'cancelled', 'error']:
|
||||||
|
self._recipe_rematch_progress = None
|
||||||
|
|
||||||
|
def is_recipe_rematch_running(self) -> bool:
|
||||||
|
"""Check if recipe rematch is currently running"""
|
||||||
|
if not self._recipe_rematch_progress:
|
||||||
|
return False
|
||||||
|
status = self._recipe_rematch_progress.get('status')
|
||||||
|
return status in ['started', 'processing']
|
||||||
|
|
||||||
def is_auto_organize_running(self) -> bool:
|
def is_auto_organize_running(self) -> bool:
|
||||||
"""Check if auto-organize is currently running"""
|
"""Check if auto-organize is currently running"""
|
||||||
if not self._auto_organize_progress:
|
if not self._auto_organize_progress:
|
||||||
|
|||||||
@@ -172,3 +172,92 @@ def test_generate_download_id(manager):
|
|||||||
download_id = manager.generate_download_id()
|
download_id = manager.generate_download_id()
|
||||||
assert isinstance(download_id, str)
|
assert isinstance(download_id, str)
|
||||||
assert download_id
|
assert download_id
|
||||||
|
|
||||||
|
|
||||||
|
# --- Recipe rematch progress channel ---
|
||||||
|
|
||||||
|
|
||||||
|
async def test_broadcast_recipe_rematch_progress_stores_and_broadcasts(manager, monkeypatch):
|
||||||
|
payload = {"status": "started", "total": 3}
|
||||||
|
broadcast_calls = []
|
||||||
|
|
||||||
|
async def fake_broadcast(data):
|
||||||
|
broadcast_calls.append(data)
|
||||||
|
|
||||||
|
monkeypatch.setattr(manager, "broadcast", fake_broadcast)
|
||||||
|
|
||||||
|
await manager.broadcast_recipe_rematch_progress(payload)
|
||||||
|
|
||||||
|
assert broadcast_calls == [payload]
|
||||||
|
assert manager.get_recipe_rematch_progress() == payload
|
||||||
|
|
||||||
|
|
||||||
|
async def test_get_recipe_rematch_progress_returns_stored(manager):
|
||||||
|
assert manager.get_recipe_rematch_progress() is None
|
||||||
|
|
||||||
|
payload = {"status": "processing", "current": 2, "total": 5}
|
||||||
|
await manager.broadcast_recipe_rematch_progress(payload)
|
||||||
|
|
||||||
|
assert manager.get_recipe_rematch_progress() == payload
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.parametrize(
|
||||||
|
"status,should_clear",
|
||||||
|
[
|
||||||
|
("started", False),
|
||||||
|
("processing", False),
|
||||||
|
("completed", True),
|
||||||
|
("cancelled", True),
|
||||||
|
("error", True),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
async def test_cleanup_recipe_rematch_progress_only_on_terminal(manager, status, should_clear):
|
||||||
|
await manager.broadcast_recipe_rematch_progress({"status": status})
|
||||||
|
|
||||||
|
manager.cleanup_recipe_rematch_progress()
|
||||||
|
|
||||||
|
if should_clear:
|
||||||
|
assert manager.get_recipe_rematch_progress() is None
|
||||||
|
else:
|
||||||
|
assert manager.get_recipe_rematch_progress() == {"status": status}
|
||||||
|
|
||||||
|
|
||||||
|
async def test_is_recipe_rematch_running_false_without_progress(manager):
|
||||||
|
assert manager.is_recipe_rematch_running() is False
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.parametrize(
|
||||||
|
"status,expected",
|
||||||
|
[
|
||||||
|
("started", True),
|
||||||
|
("processing", True),
|
||||||
|
("completed", False),
|
||||||
|
("cancelled", False),
|
||||||
|
("error", False),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
async def test_is_recipe_rematch_running_by_status(manager, status, expected):
|
||||||
|
await manager.broadcast_recipe_rematch_progress({"status": status})
|
||||||
|
|
||||||
|
assert manager.is_recipe_rematch_running() is expected
|
||||||
|
|
||||||
|
|
||||||
|
async def test_rematch_and_repair_channels_are_independent(manager):
|
||||||
|
# Rematch progress must not leak into the repair channel
|
||||||
|
await manager.broadcast_recipe_rematch_progress({"status": "processing", "current": 1})
|
||||||
|
assert manager.is_recipe_rematch_running() is True
|
||||||
|
assert manager.is_recipe_repair_running() is False
|
||||||
|
assert manager.get_recipe_repair_progress() is None
|
||||||
|
|
||||||
|
# Repair progress must not overwrite the rematch state
|
||||||
|
await manager.broadcast_recipe_repair_progress({"status": "processing", "current": 1})
|
||||||
|
assert manager.is_recipe_repair_running() is True
|
||||||
|
assert manager.is_recipe_rematch_running() is True
|
||||||
|
assert manager.get_recipe_rematch_progress() == {"status": "processing", "current": 1}
|
||||||
|
|
||||||
|
# Cleaning the rematch channel must leave the repair channel untouched
|
||||||
|
await manager.broadcast_recipe_rematch_progress({"status": "completed"})
|
||||||
|
manager.cleanup_recipe_rematch_progress()
|
||||||
|
assert manager.get_recipe_rematch_progress() is None
|
||||||
|
assert manager.get_recipe_repair_progress() == {"status": "processing", "current": 1}
|
||||||
|
assert manager.is_recipe_repair_running() is True
|
||||||
|
|||||||
Reference in New Issue
Block a user