106 lines
3.8 KiB
Python
106 lines
3.8 KiB
Python
|
|
import unittest
|
|
import sys
|
|
import os
|
|
import time
|
|
from unittest.mock import MagicMock
|
|
|
|
# --- PATH SETUP ---
|
|
# Add Project Root (ComfyUI-Doctor) to Path
|
|
current_dir = os.path.dirname(os.path.abspath(__file__))
|
|
project_root = os.path.abspath(os.path.join(current_dir, ".."))
|
|
if project_root not in sys.path:
|
|
sys.path.insert(0, project_root)
|
|
|
|
# --- MOCKING ---
|
|
mock_torch = MagicMock()
|
|
mock_torch.__version__ = "2.0.1+cu118"
|
|
mock_torch.Tensor = str
|
|
mock_torch.cuda.is_available.return_value = True
|
|
mock_torch.version.cuda = "11.8"
|
|
mock_torch.cuda.device_count.return_value = 1
|
|
mock_torch.cuda.get_device_properties.return_value.name = "Mock GPU"
|
|
mock_torch.cuda.get_device_properties.return_value.total_memory = 24 * 1024**3
|
|
sys.modules['torch'] = mock_torch
|
|
sys.modules['server'] = MagicMock() # Also mock server just in case
|
|
|
|
from analyzer import ErrorAnalyzer, NodeContext, ERROR_KEYS
|
|
from i18n import set_language, get_suggestion
|
|
from logger import SmartLogger, clear_analysis_history, get_last_analysis
|
|
|
|
class TestIntegration(unittest.TestCase):
|
|
|
|
def setUp(self):
|
|
# Ensure clean state by uninstalling any existing logger hooks
|
|
# This is needed because importing __init__ might auto-install them
|
|
SmartLogger.uninstall()
|
|
set_language("en")
|
|
|
|
def tearDown(self):
|
|
SmartLogger.uninstall()
|
|
|
|
def test_new_error_patterns(self):
|
|
"""Test the detection of the 6 new error patterns."""
|
|
|
|
test_cases = [
|
|
(
|
|
"safetensors_rust.SafetensorError: Error while deserializing header\n",
|
|
ERROR_KEYS["SAFETENSORS_ERROR"],
|
|
"SafeTensors Error"
|
|
),
|
|
(
|
|
"RuntimeError: cuDNN error: CUDNN_STATUS_EXECUTION_FAILED\n",
|
|
ERROR_KEYS["CUDNN_ERROR"],
|
|
"CUDNN Execution Failed"
|
|
),
|
|
(
|
|
"ModuleNotFoundError: No module named 'insightface'\n",
|
|
ERROR_KEYS["MISSING_INSIGHTFACE"],
|
|
"Missing InsightFace"
|
|
),
|
|
(
|
|
"RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn\n",
|
|
ERROR_KEYS["MODEL_VAE_MISMATCH"],
|
|
"Model/VAE Mismatch"
|
|
),
|
|
(
|
|
"MPS backend out of memory\n",
|
|
ERROR_KEYS["MPS_OOM"],
|
|
"MPS (Mac) OOM"
|
|
),
|
|
(
|
|
"json.decoder.JSONDecodeError: Expecting value: line 1 column 1 (char 0)\n",
|
|
ERROR_KEYS["INVALID_PROMPT"],
|
|
"Invalid Prompt Format"
|
|
)
|
|
]
|
|
|
|
for traceback_text, expected_key, expected_snippet in test_cases:
|
|
result = ErrorAnalyzer.analyze(traceback_text)
|
|
self.assertIsNotNone(result, f"Failed to detect pattern for: {expected_key}")
|
|
suggestion, _metadata = result
|
|
self.assertIsNotNone(suggestion, f"Failed to generate suggestion for: {expected_key}")
|
|
self.assertIn(expected_snippet, suggestion, f"Suggestion content mismatch for: {expected_key}")
|
|
|
|
def test_logger_validation_error_capture(self):
|
|
"""Test validation error capture via on_flush callback."""
|
|
SmartLogger.install("test.log")
|
|
clear_analysis_history()
|
|
|
|
print("Failed to validate prompt for output 1:")
|
|
print("* KSampler 1:")
|
|
print(" - Return type mismatch between linked nodes: scheduler")
|
|
print("Executing prompt: test-123")
|
|
|
|
deadline = time.time() + 2
|
|
last = get_last_analysis()
|
|
while not last.get("error") and time.time() < deadline:
|
|
time.sleep(0.05)
|
|
last = get_last_analysis()
|
|
|
|
self.assertIsNotNone(last.get("error"))
|
|
self.assertIn("Failed to validate", last.get("error", ""))
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|