Ben FinkleaandClaude Opus 4.5
9256a1c87c
A custom node extension for ComfyUI that provides: - Model version tracking for HuggingFace and CivitAI - Automatic model updates with configurable scheduling - Multi-GPU synchronization with VRAM-based filtering - Metadata management with hash verification Features: - SQLite database for model metadata storage - REST API endpoints for all operations - Web UI panel integrated into ComfyUI menu - Background scheduler for automatic update checks - rsync-based sync to remote machines - VRAM estimation for model filtering Tests: - 17 unit tests covering all library modules - Playwright integration tests for UI verification Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
360 lines
11 KiB
Python
360 lines
11 KiB
Python
"""Tests for the MetadataStore module."""
|
|
|
|
import os
|
|
import tempfile
|
|
import pytest
|
|
from pathlib import Path
|
|
|
|
import sys
|
|
sys.path.insert(0, str(Path(__file__).parent.parent))
|
|
|
|
from lib.metadata_store import MetadataStore
|
|
|
|
|
|
@pytest.fixture
|
|
def temp_db():
|
|
"""Create a temporary database for testing."""
|
|
fd, path = tempfile.mkstemp(suffix=".db")
|
|
os.close(fd)
|
|
yield path
|
|
os.unlink(path)
|
|
|
|
|
|
@pytest.fixture
|
|
def store(temp_db):
|
|
"""Create a MetadataStore instance with temporary database."""
|
|
return MetadataStore(db_path=temp_db)
|
|
|
|
|
|
class TestMetadataStore:
|
|
"""Tests for MetadataStore."""
|
|
|
|
def test_init_creates_tables(self, store):
|
|
"""Test that initialization creates required tables."""
|
|
with store._get_connection() as conn:
|
|
cursor = conn.cursor()
|
|
cursor.execute(
|
|
"SELECT name FROM sqlite_master WHERE type='table'"
|
|
)
|
|
tables = {row[0] for row in cursor.fetchall()}
|
|
|
|
assert "models" in tables
|
|
assert "sync_targets" in tables
|
|
assert "config" in tables
|
|
|
|
def test_add_and_get_model(self, store):
|
|
"""Test adding and retrieving a model."""
|
|
model_data = {
|
|
"path": "/path/to/model.safetensors",
|
|
"filename": "model.safetensors",
|
|
"model_type": "checkpoints",
|
|
"file_size": 1024 * 1024 * 1024, # 1GB
|
|
}
|
|
|
|
model_id = store.add_model(model_data)
|
|
assert model_id > 0
|
|
|
|
model = store.get_model(model_id)
|
|
assert model is not None
|
|
assert model["path"] == model_data["path"]
|
|
assert model["filename"] == model_data["filename"]
|
|
assert model["model_type"] == model_data["model_type"]
|
|
assert model["file_size"] == model_data["file_size"]
|
|
|
|
def test_get_model_by_path(self, store):
|
|
"""Test retrieving a model by path."""
|
|
model_data = {
|
|
"path": "/unique/path/model.safetensors",
|
|
"filename": "model.safetensors",
|
|
"model_type": "checkpoints",
|
|
}
|
|
|
|
store.add_model(model_data)
|
|
model = store.get_model_by_path("/unique/path/model.safetensors")
|
|
|
|
assert model is not None
|
|
assert model["path"] == model_data["path"]
|
|
|
|
def test_update_model(self, store):
|
|
"""Test updating model fields."""
|
|
model_id = store.add_model({
|
|
"path": "/path/model.safetensors",
|
|
"filename": "model.safetensors",
|
|
"model_type": "checkpoints",
|
|
})
|
|
|
|
store.update_model(model_id, {
|
|
"source_url": "https://example.com/model",
|
|
"auto_update": 1,
|
|
})
|
|
|
|
model = store.get_model(model_id)
|
|
assert model["source_url"] == "https://example.com/model"
|
|
assert model["auto_update"] == 1
|
|
|
|
def test_delete_model(self, store):
|
|
"""Test deleting a model."""
|
|
model_id = store.add_model({
|
|
"path": "/path/model.safetensors",
|
|
"filename": "model.safetensors",
|
|
"model_type": "checkpoints",
|
|
})
|
|
|
|
assert store.delete_model(model_id) is True
|
|
assert store.get_model(model_id) is None
|
|
|
|
def test_get_all_models(self, store):
|
|
"""Test retrieving all models."""
|
|
store.add_model({
|
|
"path": "/path/model1.safetensors",
|
|
"filename": "model1.safetensors",
|
|
"model_type": "checkpoints",
|
|
})
|
|
store.add_model({
|
|
"path": "/path/model2.safetensors",
|
|
"filename": "model2.safetensors",
|
|
"model_type": "loras",
|
|
})
|
|
|
|
all_models = store.get_all_models()
|
|
assert len(all_models) == 2
|
|
|
|
checkpoints = store.get_all_models(model_type="checkpoints")
|
|
assert len(checkpoints) == 1
|
|
assert checkpoints[0]["model_type"] == "checkpoints"
|
|
|
|
def test_get_models_with_updates(self, store):
|
|
"""Test getting models with updates available."""
|
|
model_id = store.add_model({
|
|
"path": "/path/model.safetensors",
|
|
"filename": "model.safetensors",
|
|
"model_type": "checkpoints",
|
|
})
|
|
|
|
# Initially no updates
|
|
assert len(store.get_models_with_updates()) == 0
|
|
|
|
# Mark as having update
|
|
store.update_model(model_id, {"update_available": 1})
|
|
updates = store.get_models_with_updates()
|
|
assert len(updates) == 1
|
|
assert updates[0]["id"] == model_id
|
|
|
|
def test_sync_target_crud(self, store):
|
|
"""Test sync target CRUD operations."""
|
|
# Add
|
|
target_id = store.add_sync_target({
|
|
"name": "test-target",
|
|
"ssh_host": "user@host",
|
|
"model_path": "/remote/path",
|
|
"max_vram_mb": 24000,
|
|
})
|
|
assert target_id > 0
|
|
|
|
# Get
|
|
target = store.get_sync_target(target_id)
|
|
assert target["name"] == "test-target"
|
|
assert target["ssh_host"] == "user@host"
|
|
assert target["max_vram_mb"] == 24000
|
|
|
|
# Update
|
|
store.update_sync_target(target_id, {"enabled": 0})
|
|
target = store.get_sync_target(target_id)
|
|
assert target["enabled"] == 0
|
|
|
|
# Delete
|
|
assert store.delete_sync_target(target_id) is True
|
|
assert store.get_sync_target(target_id) is None
|
|
|
|
def test_config_operations(self, store):
|
|
"""Test configuration get/set."""
|
|
# Default config exists
|
|
auto_update = store.get_config("auto_update_enabled")
|
|
assert auto_update == "false"
|
|
|
|
# Set and get
|
|
store.set_config("test_key", "test_value")
|
|
assert store.get_config("test_key") == "test_value"
|
|
|
|
# Get all
|
|
all_config = store.get_all_config()
|
|
assert "auto_update_enabled" in all_config
|
|
assert "test_key" in all_config
|
|
|
|
def test_get_model_count(self, store):
|
|
"""Test model count."""
|
|
assert store.get_model_count() == 0
|
|
|
|
store.add_model({
|
|
"path": "/path/model.safetensors",
|
|
"filename": "model.safetensors",
|
|
"model_type": "checkpoints",
|
|
})
|
|
|
|
assert store.get_model_count() == 1
|
|
|
|
|
|
class TestVersionTracker:
|
|
"""Tests for version tracker URL parsing."""
|
|
|
|
def test_detect_source_type(self):
|
|
"""Test source type detection from URLs."""
|
|
from lib.version_tracker import VersionTracker
|
|
tracker = VersionTracker()
|
|
|
|
assert tracker.detect_source_type(
|
|
"https://huggingface.co/org/repo/resolve/main/file.safetensors"
|
|
) == "huggingface"
|
|
|
|
assert tracker.detect_source_type(
|
|
"https://civitai.com/models/12345"
|
|
) == "civitai"
|
|
|
|
assert tracker.detect_source_type(
|
|
"https://github.com/user/repo"
|
|
) == "github"
|
|
|
|
assert tracker.detect_source_type(
|
|
"https://example.com/file.safetensors"
|
|
) == "unknown"
|
|
|
|
def test_parse_huggingface_url(self):
|
|
"""Test HuggingFace URL parsing."""
|
|
from lib.version_tracker import VersionTracker
|
|
tracker = VersionTracker()
|
|
|
|
result = tracker.parse_huggingface_url(
|
|
"https://huggingface.co/stabilityai/sdxl-vae/resolve/main/sdxl_vae.safetensors"
|
|
)
|
|
assert result is not None
|
|
repo_id, file_path = result
|
|
assert repo_id == "stabilityai/sdxl-vae"
|
|
assert file_path == "sdxl_vae.safetensors"
|
|
|
|
def test_parse_civitai_url(self):
|
|
"""Test CivitAI URL parsing."""
|
|
from lib.version_tracker import VersionTracker
|
|
tracker = VersionTracker()
|
|
|
|
# Model page
|
|
result = tracker.parse_civitai_url("https://civitai.com/models/12345")
|
|
assert result == {"model_id": "12345"}
|
|
|
|
# With version
|
|
result = tracker.parse_civitai_url(
|
|
"https://civitai.com/models/12345?modelVersionId=67890"
|
|
)
|
|
assert result == {"model_id": "12345", "version_id": "67890"}
|
|
|
|
# Download URL
|
|
result = tracker.parse_civitai_url(
|
|
"https://civitai.com/api/download/models/67890"
|
|
)
|
|
assert result == {"version_id": "67890"}
|
|
|
|
|
|
class TestSyncManager:
|
|
"""Tests for sync manager."""
|
|
|
|
def test_estimate_vram(self, temp_db):
|
|
"""Test VRAM estimation from model names."""
|
|
from lib.sync_manager import SyncManager
|
|
|
|
store = MetadataStore(db_path=temp_db)
|
|
manager = SyncManager(store)
|
|
|
|
# Test pattern matching
|
|
assert manager.estimate_vram({"filename": "wan_model.safetensors"}) == 48000
|
|
assert manager.estimate_vram({"filename": "flux-dev.safetensors"}) == 24000
|
|
assert manager.estimate_vram({"filename": "sdxl_base.safetensors"}) == 12000
|
|
assert manager.estimate_vram({"filename": "my_lora.safetensors", "model_type": "loras"}) == 2000
|
|
|
|
# Test size-based fallback
|
|
vram = manager.estimate_vram({
|
|
"filename": "unknown.safetensors",
|
|
"file_size": 5 * 1024 * 1024 * 1024 # 5GB
|
|
})
|
|
assert vram == 10000 # 5GB * 2000
|
|
|
|
def test_get_models_for_target(self, temp_db):
|
|
"""Test model filtering by VRAM."""
|
|
from lib.sync_manager import SyncManager
|
|
|
|
store = MetadataStore(db_path=temp_db)
|
|
|
|
# Add models with different VRAM requirements
|
|
store.add_model({
|
|
"path": "/path/small.safetensors",
|
|
"filename": "sd15_model.safetensors",
|
|
"model_type": "checkpoints",
|
|
"vram_mb": 6000,
|
|
})
|
|
store.add_model({
|
|
"path": "/path/large.safetensors",
|
|
"filename": "flux-dev.safetensors",
|
|
"model_type": "checkpoints",
|
|
"vram_mb": 24000,
|
|
})
|
|
|
|
manager = SyncManager(store)
|
|
|
|
# Target with 12GB limit
|
|
target = {"max_vram_mb": 12000}
|
|
models = manager.get_models_for_target(target)
|
|
|
|
assert len(models) == 1
|
|
assert models[0]["filename"] == "sd15_model.safetensors"
|
|
|
|
# Target with no limit
|
|
target = {"max_vram_mb": None}
|
|
models = manager.get_models_for_target(target)
|
|
assert len(models) == 2
|
|
|
|
|
|
class TestHashUtils:
|
|
"""Tests for hash utilities."""
|
|
|
|
def test_compute_sha256(self):
|
|
"""Test SHA256 computation."""
|
|
from lib.hash_utils import compute_sha256, verify_file_hash
|
|
|
|
# Create temp file
|
|
fd, path = tempfile.mkstemp()
|
|
try:
|
|
os.write(fd, b"test content")
|
|
os.close(fd)
|
|
|
|
hash_result = compute_sha256(path)
|
|
assert len(hash_result) == 64 # SHA256 hex length
|
|
assert hash_result.isalnum()
|
|
|
|
# Verify
|
|
assert verify_file_hash(path, hash_result) is True
|
|
assert verify_file_hash(path, "wrong_hash") is False
|
|
finally:
|
|
os.unlink(path)
|
|
|
|
def test_get_file_info(self):
|
|
"""Test file info retrieval."""
|
|
from lib.hash_utils import get_file_info
|
|
|
|
fd, path = tempfile.mkstemp()
|
|
try:
|
|
os.write(fd, b"test content")
|
|
os.close(fd)
|
|
|
|
info = get_file_info(path)
|
|
assert info["exists"] is True
|
|
assert info["size"] == 12
|
|
assert info["mtime"] is not None
|
|
finally:
|
|
os.unlink(path)
|
|
|
|
# Non-existent file
|
|
info = get_file_info("/nonexistent/file.txt")
|
|
assert info["exists"] is False
|
|
|
|
|
|
if __name__ == "__main__":
|
|
pytest.main([__file__, "-v"])
|