mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-03-24 14:42:11 -03:00
test(routes): tidy lora route test imports
This commit is contained in:
84
tests/services/test_websocket_manager.py
Normal file
84
tests/services/test_websocket_manager.py
Normal file
@@ -0,0 +1,84 @@
|
||||
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_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
|
||||
Reference in New Issue
Block a user