/**
* Test script for VRAM cleanup functionality
* Run: node test/test-vram-cleanup.js
*/
import { SimpleVRAMManager } from '../src/services/simple-vram-manager.js';
// Mock ComfyUI client for testing
class MockComfyUIClient {
constructor() {
this.serverAddress = '127.0.0.1:8188';
this.promptCount = 0;
}
async queuePrompt(workflow) {
console.log('📤 Queuing workflow:', Object.keys(workflow).length, 'nodes');
this.promptCount++;
return { prompt_id: `test-${this.promptCount}` };
}
async waitForExecution(promptId) {
console.log(`⏳ Waiting for ${promptId}...`);
await new Promise(resolve => setTimeout(resolve, 1000));
console.log(`✅ ${promptId} completed`);
}
generateImage(params) {
console.log('🎨 Generating image:', params.model || 'default');
return Promise.resolve({ success: true });
}
upscaleImage(params) {
console.log('🔍 Upscaling image');
return Promise.resolve({ success: true });
}
}
async function testVRAMCleanup() {
console.log('🧪 Testing VRAM Cleanup System\n');
console.log('=' .repeat(50));
// Create mock client
const client = new MockComfyUIClient();
// Create VRAM manager with test configuration
const vramManager = new SimpleVRAMManager(client, {
warningThreshold: 0.75,
criticalThreshold: 0.85,
cleanupThreshold: 0.90,
checkInterval: 5000, // Check every 5 seconds for testing
idleTimeout: 10000, // 10 seconds idle for testing
cooldownPeriod: 5000, // 5 second cooldown for testing
autoCleanup: true,
preRequestCheck: true,
logging: true
});
console.log('\n📊 Test 1: Status Check');
console.log('-'.repeat(30));
const status = vramManager.getStatus();
console.log('Initial status:', JSON.stringify(status, null, 2));
console.log('\n🔄 Test 2: Manual VRAM Check');
console.log('-'.repeat(30));
const checkResult = await vramManager.checkAndCleanup();
console.log('Check result:', checkResult);
console.log('\n🧹 Test 3: Force Cleanup');
console.log('-'.repeat(30));
const cleanupResult = await vramManager.forceCleanup();
console.log('Cleanup result:', cleanupResult);
console.log('\n🚀 Test 4: Pre-Request Check (FLUX)');
console.log('-'.repeat(30));
const preCheckFlux = await vramManager.preRequestCheck('flux');
console.log('Pre-check for FLUX:', preCheckFlux);
console.log('\n🖼️ Test 5: Pre-Request Check (Upscale)');
console.log('-'.repeat(30));
const preCheckUpscale = await vramManager.preRequestCheck('upscale');
console.log('Pre-check for upscale:', preCheckUpscale);
console.log('\n⏱️ Test 6: Idle Detection');
console.log('-'.repeat(30));
vramManager.updateActivity();
console.log('Activity updated, idle:', vramManager.isIdle());
console.log('Waiting 11 seconds for idle timeout...');
await new Promise(resolve => setTimeout(resolve, 11000));
console.log('After wait, idle:', vramManager.isIdle());
console.log('\n🔧 Test 7: Setup with Hooks');
console.log('-'.repeat(30));
const { setupVRAMManagement } = await import('../src/services/simple-vram-manager.js');
const managedClient = new MockComfyUIClient();
const manager = setupVRAMManagement(managedClient);
// Test hooked generate
await managedClient.generateImage({ model: 'flux', prompt: 'test' });
// Test hooked upscale
await managedClient.upscaleImage({ image_path: 'test.png' });
console.log('\n✅ Test 8: Monitoring Start/Stop');
console.log('-'.repeat(30));
manager.start();
console.log('Monitoring started');
await new Promise(resolve => setTimeout(resolve, 2000));
manager.stop();
console.log('Monitoring stopped');
console.log('\n' + '='.repeat(50));
console.log('🎉 All tests completed!');
// Cleanup
process.exit(0);
}
// Run tests
testVRAMCleanup().catch(error => {
console.error('❌ Test failed:', error);
process.exit(1);
});
