mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-03-22 13:42:12 -03:00
104 lines
3.0 KiB
Python
104 lines
3.0 KiB
Python
from datetime import datetime, timedelta
|
|
|
|
import pytest
|
|
|
|
from py.services.websocket_manager import WebSocketManager
|
|
|
|
|
|
class DummyWebSocket:
|
|
def __init__(self):
|
|
self.messages = []
|
|
self.closed = False
|
|
|
|
async def send_json(self, data):
|
|
if self.closed:
|
|
raise RuntimeError("WebSocket closed")
|
|
self.messages.append(data)
|
|
|
|
|
|
@pytest.fixture
|
|
def manager():
|
|
return WebSocketManager()
|
|
|
|
|
|
async def test_broadcast_init_progress_adds_defaults(manager):
|
|
ws = DummyWebSocket()
|
|
manager._init_websockets.add(ws)
|
|
|
|
await manager.broadcast_init_progress({})
|
|
|
|
assert ws.messages == [
|
|
{
|
|
"stage": "processing",
|
|
"progress": 0,
|
|
"details": "Processing...",
|
|
}
|
|
]
|
|
|
|
|
|
async def test_broadcast_init_progress_caches_payload(manager):
|
|
await manager.broadcast_init_progress({'pageType': 'loras', 'progress': 42})
|
|
|
|
cached = manager._last_init_progress.get('page:loras')
|
|
assert cached is not None
|
|
assert cached['progress'] == 42
|
|
assert cached['stage'] == 'processing'
|
|
assert cached['details'] == 'Processing...'
|
|
|
|
|
|
async def test_send_cached_progress_to_new_client(manager):
|
|
await manager.broadcast_init_progress({'pageType': 'loras', 'progress': 87})
|
|
|
|
ws = DummyWebSocket()
|
|
await manager._send_cached_init_progress(ws)
|
|
|
|
assert ws.messages[-1]['progress'] == 87
|
|
assert ws.messages[-1]['pageType'] == 'loras'
|
|
|
|
async def test_broadcast_download_progress_tracks_state(manager):
|
|
ws = DummyWebSocket()
|
|
download_id = "abc"
|
|
manager._download_websockets[download_id] = ws
|
|
|
|
await manager.broadcast_download_progress(download_id, {"progress": 55})
|
|
|
|
assert ws.messages == [{"progress": 55}]
|
|
assert manager.get_download_progress(download_id)["progress"] == 55
|
|
|
|
|
|
async def test_broadcast_download_progress_missing_socket(manager):
|
|
await manager.broadcast_download_progress("missing", {"progress": 30})
|
|
# Progress should be stored even without a live websocket
|
|
assert manager.get_download_progress("missing")["progress"] == 30
|
|
|
|
|
|
async def test_auto_organize_progress_helpers(manager):
|
|
payload = {"status": "processing", "progress": 10}
|
|
await manager.broadcast_auto_organize_progress(payload)
|
|
|
|
assert manager.get_auto_organize_progress() == payload
|
|
assert manager.is_auto_organize_running() is True
|
|
|
|
manager.cleanup_auto_organize_progress()
|
|
assert manager.get_auto_organize_progress() is None
|
|
assert manager.is_auto_organize_running() is False
|
|
|
|
|
|
def test_cleanup_old_downloads(manager):
|
|
now = datetime.now()
|
|
manager._download_progress = {
|
|
"recent": {"progress": 10, "timestamp": now},
|
|
"stale": {"progress": 100, "timestamp": now - timedelta(hours=48)},
|
|
}
|
|
|
|
manager.cleanup_old_downloads(max_age_hours=24)
|
|
|
|
assert "stale" not in manager._download_progress
|
|
assert "recent" in manager._download_progress
|
|
|
|
|
|
def test_generate_download_id(manager):
|
|
download_id = manager.generate_download_id()
|
|
assert isinstance(download_id, str)
|
|
assert download_id
|