352 lines
12 KiB
Python
352 lines
12 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Test script for SageUtils custom routes.
|
|
This demonstrates how the custom routes would work when ComfyUI is running.
|
|
"""
|
|
|
|
import sys
|
|
import os
|
|
import json
|
|
from typing import Dict, Any
|
|
from unittest.mock import Mock, MagicMock
|
|
|
|
# Add ComfyUI path
|
|
sys.path.insert(0, '/home/ai/programs/comfyui')
|
|
|
|
class MockRequest:
|
|
"""Mock aiohttp request object"""
|
|
def __init__(self, match_info=None):
|
|
self.match_info = match_info or {}
|
|
|
|
class MockWeb:
|
|
"""Mock aiohttp web module"""
|
|
@staticmethod
|
|
def json_response(data, status=200):
|
|
return {
|
|
"status": status,
|
|
"content_type": "application/json",
|
|
"data": data
|
|
}
|
|
|
|
class MockRoutes:
|
|
"""Mock routes object that stores registered routes"""
|
|
def __init__(self):
|
|
self.routes = {}
|
|
|
|
def get(self, path):
|
|
def decorator(func):
|
|
self.routes[f"GET {path}"] = func
|
|
print(f"✓ Registered route: GET {path}")
|
|
return func
|
|
return decorator
|
|
|
|
async def call_route(self, method, path, request=None):
|
|
"""Call a registered route for testing"""
|
|
route_key = f"{method} {path}"
|
|
if route_key in self.routes:
|
|
return await self.routes[route_key](request or MockRequest())
|
|
else:
|
|
return MockWeb.json_response({"error": "Route not found"}, 404)
|
|
|
|
class MockPromptServer:
|
|
"""Mock PromptServer class"""
|
|
def __init__(self):
|
|
self.routes = MockRoutes()
|
|
|
|
@classmethod
|
|
def create_instance(cls):
|
|
cls.instance = cls()
|
|
return cls.instance
|
|
|
|
def test_sage_routes():
|
|
"""Test the SageUtils custom routes"""
|
|
print("=== Testing SageUtils Custom Routes ===\n")
|
|
|
|
# Create mock server instance
|
|
prompt_server = MockPromptServer.create_instance()
|
|
|
|
# Mock the modules that would be imported
|
|
sys.modules['server'] = Mock()
|
|
sys.modules['server'].PromptServer = MockPromptServer
|
|
sys.modules['aiohttp'] = Mock()
|
|
sys.modules['aiohttp'].web = MockWeb
|
|
|
|
# Import and setup our cache
|
|
print("1. Loading SageUtils cache...")
|
|
try:
|
|
# Suppress hash warnings
|
|
import warnings
|
|
warnings.filterwarnings("ignore")
|
|
|
|
from custom_nodes.comfyui_sageutils.utils.model_cache import cache
|
|
cache.load()
|
|
print(f" ✓ Cache loaded: {len(cache.hash)} files, {len(cache.info)} info entries")
|
|
|
|
# Add some dummy data for testing if cache is empty
|
|
if len(cache.hash) == 0:
|
|
print(" Adding dummy test data...")
|
|
test_hash = "abc123def456789"
|
|
test_path = "/models/test_model.safetensors"
|
|
cache.hash[test_path] = test_hash
|
|
cache.info[test_hash] = {
|
|
"hash": test_hash,
|
|
"civitai": "true",
|
|
"modelId": "12345",
|
|
"name": "Test Model",
|
|
"baseModel": "SDXL",
|
|
"lastUsed": "2025-06-22T10:30:00"
|
|
}
|
|
print(f" ✓ Added test data: 1 file, 1 info entry")
|
|
|
|
except Exception as e:
|
|
print(f" ✗ Error loading cache: {e}")
|
|
import traceback
|
|
traceback.print_exc()
|
|
return False
|
|
|
|
# Mock the imports for our routes module
|
|
original_modules = {}
|
|
mock_modules = {
|
|
'server': Mock(),
|
|
'aiohttp': Mock(),
|
|
'.utils.model_cache': Mock()
|
|
}
|
|
|
|
# Setup mocks
|
|
mock_modules['server'].PromptServer = MockPromptServer
|
|
mock_modules['aiohttp'].web = MockWeb
|
|
mock_modules['.utils.model_cache'].cache = cache
|
|
|
|
print("\n2. Testing route registration...")
|
|
|
|
# Manually register routes like our module would
|
|
routes = prompt_server.routes
|
|
|
|
@routes.get('/sage_cache/info')
|
|
async def get_sage_cache_info(request):
|
|
try:
|
|
cache.load()
|
|
return MockWeb.json_response(cache.info)
|
|
except Exception as e:
|
|
return MockWeb.json_response(
|
|
{"error": f"Failed to retrieve cache info: {str(e)}"},
|
|
status=500
|
|
)
|
|
|
|
@routes.get('/sage_cache/hash')
|
|
async def get_sage_cache_hash(request):
|
|
try:
|
|
cache.load()
|
|
return MockWeb.json_response(cache.hash)
|
|
except Exception as e:
|
|
return MockWeb.json_response(
|
|
{"error": f"Failed to retrieve cache hash: {str(e)}"},
|
|
status=500
|
|
)
|
|
|
|
@routes.get('/sage_cache/stats')
|
|
async def get_sage_cache_stats(request):
|
|
try:
|
|
cache.load()
|
|
|
|
stats = {
|
|
"total_files": len(cache.hash),
|
|
"total_info_entries": len(cache.info),
|
|
"cache_files": {
|
|
"hash_path": str(cache.hash_path),
|
|
"info_path": str(cache.info_path),
|
|
"main_path": str(cache.main_path),
|
|
"ollama_models_path": str(cache.ollama_models_path)
|
|
}
|
|
}
|
|
|
|
# Count files by civitai status
|
|
civitai_found = 0
|
|
civitai_not_found = 0
|
|
|
|
for info_data in cache.info.values():
|
|
civitai_status = info_data.get("civitai", "False")
|
|
if isinstance(civitai_status, str):
|
|
if civitai_status.lower() in ["true", "1"]:
|
|
civitai_found += 1
|
|
else:
|
|
civitai_not_found += 1
|
|
elif civitai_status is True:
|
|
civitai_found += 1
|
|
else:
|
|
civitai_not_found += 1
|
|
|
|
stats["civitai_stats"] = {
|
|
"found_on_civitai": civitai_found,
|
|
"not_found_on_civitai": civitai_not_found
|
|
}
|
|
|
|
return MockWeb.json_response(stats)
|
|
except Exception as e:
|
|
return MockWeb.json_response(
|
|
{"error": f"Failed to retrieve cache stats: {str(e)}"},
|
|
status=500
|
|
)
|
|
|
|
@routes.get('/sage_cache/file/{file_hash}')
|
|
async def get_sage_cache_file_info(request):
|
|
try:
|
|
file_hash = request.match_info.get('file_hash', '')
|
|
if not file_hash:
|
|
return MockWeb.json_response(
|
|
{"error": "No file hash provided"},
|
|
status=400
|
|
)
|
|
|
|
cache.load()
|
|
|
|
file_info = cache.info.get(file_hash)
|
|
if file_info is None:
|
|
return MockWeb.json_response(
|
|
{"error": f"No information found for hash: {file_hash}"},
|
|
status=404
|
|
)
|
|
|
|
# Also include which file paths use this hash
|
|
file_paths = [path for path, hash_val in cache.hash.items() if hash_val == file_hash]
|
|
|
|
result = {
|
|
"hash": file_hash,
|
|
"info": file_info,
|
|
"file_paths": file_paths
|
|
}
|
|
|
|
return MockWeb.json_response(result)
|
|
except Exception as e:
|
|
return MockWeb.json_response(
|
|
{"error": f"Failed to retrieve file info: {str(e)}"},
|
|
status=500
|
|
)
|
|
|
|
@routes.get('/sage_cache/path')
|
|
async def get_sage_cache_path_info(request):
|
|
try:
|
|
# Get file path from query parameter (mocked)
|
|
file_path = getattr(request, 'file_path', '/models/test_model.safetensors')
|
|
if not file_path:
|
|
return MockWeb.json_response(
|
|
{"error": "No file_path parameter provided"},
|
|
status=400
|
|
)
|
|
|
|
cache.load()
|
|
|
|
# Look up hash for this file path
|
|
file_hash = cache.hash.get(file_path)
|
|
if file_hash is None:
|
|
return MockWeb.json_response(
|
|
{"error": f"No hash found for file path: {file_path}"},
|
|
status=404
|
|
)
|
|
|
|
# Get the info for this hash
|
|
file_info = cache.info.get(file_hash)
|
|
if file_info is None:
|
|
return MockWeb.json_response(
|
|
{"error": f"No information found for hash {file_hash} (path: {file_path})"},
|
|
status=404
|
|
)
|
|
|
|
# Also include all file paths that use this same hash
|
|
all_file_paths = [path for path, hash_val in cache.hash.items() if hash_val == file_hash]
|
|
|
|
result = {
|
|
"file_path": file_path,
|
|
"hash": file_hash,
|
|
"info": file_info,
|
|
"all_paths_with_same_hash": all_file_paths
|
|
}
|
|
|
|
return MockWeb.json_response(result)
|
|
except Exception as e:
|
|
return MockWeb.json_response(
|
|
{"error": f"Failed to retrieve path info: {str(e)}"},
|
|
status=500
|
|
)
|
|
|
|
print("\n3. Testing routes...")
|
|
|
|
# Test each route
|
|
async def test_routes():
|
|
# Test /sage_cache/stats
|
|
print("\n Testing GET /sage_cache/stats")
|
|
response = await routes.call_route("GET", "/sage_cache/stats")
|
|
print(f" Status: {response['status']}")
|
|
if response['status'] == 200:
|
|
data = response['data']
|
|
print(f" Total files: {data['total_files']}")
|
|
print(f" Total info entries: {data['total_info_entries']}")
|
|
print(f" Civitai found: {data['civitai_stats']['found_on_civitai']}")
|
|
print(f" Civitai not found: {data['civitai_stats']['not_found_on_civitai']}")
|
|
|
|
# Test /sage_cache/hash
|
|
print("\n Testing GET /sage_cache/hash")
|
|
response = await routes.call_route("GET", "/sage_cache/hash")
|
|
print(f" Status: {response['status']}")
|
|
if response['status'] == 200:
|
|
data = response['data']
|
|
print(f" Hash entries: {len(data)}")
|
|
if data:
|
|
first_path = list(data.keys())[0]
|
|
print(f" Sample: {first_path[:50]}... -> {data[first_path][:16]}...")
|
|
|
|
# Test /sage_cache/info
|
|
print("\n Testing GET /sage_cache/info")
|
|
response = await routes.call_route("GET", "/sage_cache/info")
|
|
print(f" Status: {response['status']}")
|
|
if response['status'] == 200:
|
|
data = response['data']
|
|
print(f" Info entries: {len(data)}")
|
|
if data:
|
|
first_hash = list(data.keys())[0]
|
|
print(f" Sample hash: {first_hash}")
|
|
info = data[first_hash]
|
|
print(f" Sample info: civitai={info.get('civitai')}, name={info.get('name', 'N/A')}")
|
|
|
|
# Test /sage_cache/file/{hash}
|
|
if cache.info:
|
|
test_hash = list(cache.info.keys())[0]
|
|
print(f"\n Testing GET /sage_cache/file/{test_hash}")
|
|
request = MockRequest({"file_hash": test_hash})
|
|
response = await routes.call_route("GET", "/sage_cache/file/{file_hash}", request)
|
|
print(f" Status: {response['status']}")
|
|
if response['status'] == 200:
|
|
data = response['data']
|
|
print(f" Hash: {data['hash']}")
|
|
print(f" File paths: {len(data['file_paths'])}")
|
|
print(f" Info keys: {list(data['info'].keys())}")
|
|
|
|
# Test /sage_cache/path
|
|
print("\n Testing GET /sage_cache/path")
|
|
response = await routes.call_route("GET", "/sage_cache/path")
|
|
print(f" Status: {response['status']}")
|
|
if response['status'] == 200:
|
|
data = response['data']
|
|
print(f" File path: {data['file_path']}")
|
|
print(f" Hash: {data['hash']}")
|
|
print(f" Info keys: {list(data['info'].keys())}")
|
|
print(f" All paths with same hash: {len(data['all_paths_with_same_hash'])}")
|
|
|
|
# Run the async tests
|
|
import asyncio
|
|
asyncio.run(test_routes())
|
|
|
|
print("\n=== Test Complete ===")
|
|
print("✓ All routes registered and tested successfully!")
|
|
print("\nWhen ComfyUI is running, these routes will be available at:")
|
|
print(" - http://localhost:8188/sage_cache/info")
|
|
print(" - http://localhost:8188/sage_cache/hash")
|
|
print(" - http://localhost:8188/sage_cache/stats")
|
|
print(" - http://localhost:8188/sage_cache/file/{hash}")
|
|
print(" - Or with /api prefix for each route")
|
|
|
|
return True
|
|
|
|
if __name__ == "__main__":
|
|
test_sage_routes()
|