tests/verify-injection.js

import { router } from '../src/endpoints/secure-generate.js';
import express from 'express';
import fetch from 'node-fetch';
import fs from 'node:fs';
import path from 'node:path';
import { fileURLToPath } from 'url';

const __dirname = path.dirname(fileURLToPath(import.meta.url));
const rootPath = path.join(__dirname, '..');

// Mock InternGen global
globalThis.DATA_ROOT = path.join(rootPath, 'data');

// Create a dummy hidden_prompts.json if it doesn't exist (it should from my previous step)
const testPrompts = {
    "test_id": {
        "name": "Test Prompt",
        "prompt": "INJECTED_PROMPT_CONTENT",
        "role": "system"
    }
};
fs.writeFileSync(path.join(globalThis.DATA_ROOT, 'hidden_prompts.json'), JSON.stringify(testPrompts));

const app = express();
app.use(express.json());
app.use('/test', router);

// Start a mock upstream server
const upstreamApp = express();
upstreamApp.use(express.json());
upstreamApp.post('/v1/chat/completions', (req, res) => {
    console.log('Upstream received messages:', JSON.stringify(req.body.messages, null, 2));
    const injected = req.body.messages.some(m => m.content === 'INJECTED_PROMPT_CONTENT');
    res.json({ success: true, injected: injected });
});

const upstreamServer = upstreamApp.listen(5001, async () => {
    const server = app.listen(5002, async () => {
        try {
            console.log('Running injection test...');

            // Test /list endpoint
            const listResponse = await fetch('http://localhost:5002/test/list');
            const listResult = await listResponse.json();
            console.log('List Result:', listResult);
            if (listResult.some(p => p.id === 'test_id' && p.label === 'Test Prompt')) {
                console.log('VERIFICATION SUCCESS: /list endpoint returned correct buttons!');
            } else {
                console.log('VERIFICATION FAILED: /list endpoint returned incorrect data.');
                process.exit(1);
            }

            // Note: The previous call was to 5001 directly to test upstream. 
            // Now let's call 5002 which is our proxy.
            const proxyResponse = await fetch('http://localhost:5002/test', {
                method: 'POST',
                headers: { 'Content-Type': 'application/json' },
                body: JSON.stringify({
                    target_url: 'http://localhost:5001/v1/chat/completions',
                    hidden_prompt_id: 'test_id',
                    messages: [{ role: 'user', content: 'hello' }]
                })
            });

            const result = await proxyResponse.json();
            console.log('Result:', result);

            if (result.injected) {
                console.log('VERIFICATION SUCCESS: Prompt injected correctly!');
            } else {
                console.log('VERIFICATION FAILED: Prompt not found in upstream request.');
                process.exit(1);
            }
        } catch (err) {
            console.error('Test error:', err);
            process.exit(1);
        } finally {
            server.close();
            upstreamServer.close();
            process.exit(0);
        }
    });
});