#!/usr/bin/env node
/**
* Test script for connection resilience features
* Tests automatic reconnection, exponential backoff, and operation queuing
*/
import { ComfyUIClient } from '../src/comfyui-client.js';
import { getFluxDiffusersWorkflow } from '../src/flux-workflow.js';
// ANSI color codes for output
const colors = {
reset: '\x1b[0m',
red: '\x1b[31m',
green: '\x1b[32m',
yellow: '\x1b[33m',
blue: '\x1b[34m',
cyan: '\x1b[36m'
};
function log(message, color = 'reset') {
console.log(`${colors[color]}${message}${colors.reset}`);
}
async function testBasicConnection() {
log('\n=== Test 1: Basic Connection ===', 'cyan');
const client = new ComfyUIClient('comfyui:8188', {
maxReconnectAttempts: 3,
reconnectDelay: 500,
heartbeatInterval: 5000
});
try {
await client.connect();
log('✅ Successfully connected', 'green');
const status = client.getConnectionStatus();
log(`Connection status: ${JSON.stringify(status, null, 2)}`, 'blue');
client.disconnect();
log('✅ Successfully disconnected', 'green');
return true;
} catch (error) {
log(`❌ Connection failed: ${error.message}`, 'red');
return false;
}
}
async function testReconnection() {
log('\n=== Test 2: Automatic Reconnection ===', 'cyan');
log('This test simulates connection loss and recovery', 'yellow');
const client = new ComfyUIClient('comfyui:8188', {
maxReconnectAttempts: 5,
reconnectDelay: 1000,
heartbeatInterval: 3000
});
try {
await client.connect();
log('✅ Initial connection established', 'green');
// Listen for reconnection events
client.connectionManager.on('stateChange', ({ from, to }) => {
log(`State change: ${from} -> ${to}`, 'yellow');
});
// Simulate connection loss by closing WebSocket
log('Simulating connection loss...', 'yellow');
if (client.ws) {
client.ws.close();
}
// Wait for automatic reconnection
await new Promise(resolve => {
client.connectionManager.once('connected', () => {
log('✅ Successfully reconnected!', 'green');
resolve();
});
// Timeout after 10 seconds
setTimeout(() => {
log('⚠️ Reconnection timeout', 'yellow');
resolve();
}, 10000);
});
client.disconnect();
return true;
} catch (error) {
log(`❌ Test failed: ${error.message}`, 'red');
return false;
}
}
async function testOperationQueuing() {
log('\n=== Test 3: Operation Queuing During Reconnection ===', 'cyan');
const client = new ComfyUIClient('comfyui:8188', {
maxReconnectAttempts: 5,
reconnectDelay: 500
});
try {
await client.connect();
log('✅ Connected', 'green');
// Queue an operation
const workflow = getFluxDiffusersWorkflow({
prompt: 'Test image for resilience testing',
width: 512,
height: 512,
steps: 1,
cfg_scale: 1.0,
seed: 42
});
log('Queueing operation...', 'yellow');
// This should work even if connection is temporarily lost
const promise = client.generateImage(workflow, './output');
// Simulate brief connection loss
if (client.ws) {
log('Simulating brief connection loss...', 'yellow');
client.ws.close();
}
// The operation should complete after reconnection
const result = await promise;
if (result && result.length > 0) {
log(`✅ Operation completed successfully! Generated ${result.length} image(s)`, 'green');
} else {
log('⚠️ Operation completed but no images generated', 'yellow');
}
client.disconnect();
return true;
} catch (error) {
log(`❌ Test failed: ${error.message}`, 'red');
return false;
}
}
async function testExponentialBackoff() {
log('\n=== Test 4: Exponential Backoff ===', 'cyan');
// Test with intentionally wrong address to trigger backoff
const client = new ComfyUIClient('invalid-host:8188', {
maxReconnectAttempts: 3,
reconnectDelay: 100,
maxReconnectDelay: 1000
});
const delays = [];
let lastAttemptTime = Date.now();
client.connectionManager.on('stateChange', ({ to }) => {
if (to === 'reconnecting') {
const now = Date.now();
const delay = now - lastAttemptTime;
lastAttemptTime = now;
delays.push(delay);
log(`Reconnection attempt after ${delay}ms`, 'yellow');
}
});
try {
await client.connect();
log('⚠️ Unexpected success', 'yellow');
} catch (error) {
log('✅ Connection failed as expected', 'green');
// Check if delays are increasing (exponential backoff)
let isExponential = true;
for (let i = 1; i < delays.length; i++) {
if (delays[i] <= delays[i - 1]) {
isExponential = false;
break;
}
}
if (isExponential && delays.length > 1) {
log('✅ Exponential backoff working correctly', 'green');
log(`Delays: ${delays.join('ms, ')}ms`, 'blue');
return true;
} else {
log('❌ Exponential backoff not working correctly', 'red');
return false;
}
}
}
async function testConnectionStatus() {
log('\n=== Test 5: Connection Status Monitoring ===', 'cyan');
const client = new ComfyUIClient('comfyui:8188');
try {
log('Initial status:', 'yellow');
log(JSON.stringify(client.getConnectionStatus(), null, 2), 'blue');
await client.connect();
log('\nConnected status:', 'yellow');
log(JSON.stringify(client.getConnectionStatus(), null, 2), 'blue');
client.disconnect();
log('\nDisconnected status:', 'yellow');
log(JSON.stringify(client.getConnectionStatus(), null, 2), 'blue');
log('✅ Status monitoring working correctly', 'green');
return true;
} catch (error) {
log(`❌ Test failed: ${error.message}`, 'red');
return false;
}
}
async function runAllTests() {
log('🧪 Starting Connection Resilience Tests', 'cyan');
log('=' .repeat(50), 'cyan');
const tests = [
{ name: 'Basic Connection', fn: testBasicConnection },
{ name: 'Automatic Reconnection', fn: testReconnection },
{ name: 'Operation Queuing', fn: testOperationQueuing },
{ name: 'Exponential Backoff', fn: testExponentialBackoff },
{ name: 'Connection Status', fn: testConnectionStatus }
];
const results = [];
for (const test of tests) {
try {
const passed = await test.fn();
results.push({ name: test.name, passed });
} catch (error) {
log(`❌ Test "${test.name}" threw an error: ${error.message}`, 'red');
results.push({ name: test.name, passed: false });
}
}
// Print summary
log('\n' + '=' .repeat(50), 'cyan');
log('📊 Test Summary', 'cyan');
log('=' .repeat(50), 'cyan');
let passedCount = 0;
for (const result of results) {
const icon = result.passed ? '✅' : '❌';
const color = result.passed ? 'green' : 'red';
log(`${icon} ${result.name}`, color);
if (result.passed) passedCount++;
}
log('\n' + '=' .repeat(50), 'cyan');
const allPassed = passedCount === results.length;
const summaryColor = allPassed ? 'green' : 'yellow';
log(`Results: ${passedCount}/${results.length} tests passed`, summaryColor);
if (allPassed) {
log('🎉 All tests passed!', 'green');
} else {
log('⚠️ Some tests failed. Please review the output above.', 'yellow');
}
process.exit(allPassed ? 0 : 1);
}
// Run tests
runAllTests().catch(error => {
log(`Fatal error: ${error.message}`, 'red');
process.exit(1);
});
