test.js

#!/usr/bin/env node

// Simple test for the Z.AI proxy worker
const workerModule = await import('./index.js');
const worker = workerModule.default;

// Mock environment
const env = {
  TOKEN_POOL_SIZE: process.env.TOKEN_POOL_SIZE || '2',
  API_KEY: process.env.API_KEY || 'sk-z2api-key-2024',
  SHOW_THINK_TAGS: process.env.SHOW_THINK_TAGS || 'true'
};

async function testHealthEndpoint() {
  const request = new Request('http://localhost:8787/health');
  const response = await worker.fetch(request, env);
  const result = await response.json();
  console.log('✓ Health check:', result);
  return response.status === 200;
}

async function testModelsEndpoint() {
  const request = new Request('http://localhost:8787/v1/models');
  const response = await worker.fetch(request, env);
  const result = await response.json();
  console.log('✓ Models endpoint:', result);
  return response.status === 200 && result.data.length > 0;
}

async function testChatCompletion() {
  try {
    const request = new Request('http://localhost:8787/v1/chat/completions', {
      method: 'POST',
      headers: {
        'Content-Type': 'application/json',
        'Authorization': 'Bearer sk-z2api-key-2024'
      },
      body: JSON.stringify({
        model: 'glm-4.5-air',
        messages: [{ role: 'user', content: 'Say hello' }],
        stream: false
      })
    });
    
    const response = await worker.fetch(request, env);
    console.log('✓ Chat completion status:', response.status);
    
    if (response.status === 200) {
      const result = await response.json();
      console.log('✓ Full chat response:', JSON.stringify(result, null, 2));
      if (result.choices && result.choices[0] && result.choices[0].message) {
        console.log('✓ Message content:', result.choices[0].message.content);
        return result.choices[0].message.content.length > 0;
      } else {
        console.log('✗ No valid message content in response');
        return false;
      }
    } else if (response.status === 503) {
      console.log('⚠️  Token pool initialization failed - chat test not applicable');
      throw new Error('Token pool required for chat test');
    } else if (response.status === 500) {
      const error = await response.text();
      if (error.includes('Internal server error')) {
        console.log('⚠️  Token fetch likely failed - network or service issue');
        throw new Error('Token pool required for chat test');
      } else {
        console.log('✗ Chat completion error:', error);
        return false;
      }
    } else {
      const error = await response.text();
      console.log('✗ Chat completion error:', error);
      return false;
    }
  } catch (error) {
    console.log('⚠️  Chat test failed:', error.message);
    throw error;
  }
}

async function testStreamingChat() {
  try {
    const request = new Request('http://localhost:8787/v1/chat/completions', {
      method: 'POST',
      headers: {
        'Content-Type': 'application/json',
        'Authorization': 'Bearer sk-z2api-key-2024'
      },
      body: JSON.stringify({
        model: 'glm-4.5-air',
        messages: [{ role: 'user', content: 'Count to 3' }],
        stream: true
      })
    });
    
    const response = await worker.fetch(request, env);
    console.log('✓ Streaming chat status:', response.status);
    
    if (response.status === 200) {
      const reader = response.body.getReader();
      const decoder = new TextDecoder();
      let chunks = 0;
      
      try {
        while (true) {
          const { done, value } = await reader.read();
          if (done) break;
          
          const chunk = decoder.decode(value);
          chunks++;
          if (chunks <= 3) {
            console.log(`✓ Stream chunk ${chunks}:`, chunk.substring(0, 100) + '...');
          }
          
          if (chunks > 20) break; // Prevent infinite loop
        }
        
        console.log(`✓ Received ${chunks} streaming chunks`);
        return chunks > 0;
      } catch (error) {
        console.log('✗ Streaming error:', error.message);
        return false;
      }
    } else if (response.status === 503) {
      console.log('⚠️  Token pool initialization failed - streaming test not applicable');
      throw new Error('Token pool required for streaming test');
    } else if (response.status === 500) {
      const error = await response.text();
      if (error.includes('Internal server error')) {
        console.log('⚠️  Token fetch likely failed - network or service issue');
        throw new Error('Token pool required for streaming test');
      } else {
        console.log('✗ Streaming chat error:', error);
        return false;
      }
    } else {
      const error = await response.text();
      console.log('✗ Streaming chat error:', error);
      return false;
    }
  } catch (error) {
    console.log('⚠️  Streaming test failed:', error.message);
    throw error;
  }
}

async function testInvalidAuth() {
  const request = new Request('http://localhost:8787/v1/chat/completions', {
    method: 'POST',
    headers: {
      'Content-Type': 'application/json',
      'Authorization': 'Bearer invalid-key'
    },
    body: JSON.stringify({
      model: 'glm-4.5-air',
      messages: [{ role: 'user', content: 'Hello' }],
      stream: false
    })
  });
  
  const response = await worker.fetch(request, env);
  console.log('✓ Invalid auth status:', response.status);
  return response.status === 401;
}

async function testInvalidModel() {
  const request = new Request('http://localhost:8787/v1/chat/completions', {
    method: 'POST',
    headers: {
      'Content-Type': 'application/json',
      'Authorization': 'Bearer sk-z2api-key-2024'
    },
    body: JSON.stringify({
      model: 'invalid-model',
      messages: [{ role: 'user', content: 'Hello' }],
      stream: false
    })
  });
  
  const response = await worker.fetch(request, env);
  console.log('✓ Invalid model status:', response.status);
  return response.status === 404;
}

async function testEmptyTokenPool() {
  // Test with zero token pool size - need a fresh worker instance
  const freshWorkerModule = await import(`./index.js?t=${Date.now()}`);
  const freshWorker = freshWorkerModule.default;
  const emptyPoolEnv = { TOKEN_POOL_SIZE: '0' };
  
  const request = new Request('http://localhost:8787/v1/chat/completions', {
    method: 'POST',
    headers: {
      'Content-Type': 'application/json',
      'Authorization': 'Bearer sk-z2api-key-2024'
    },
    body: JSON.stringify({
      model: 'glm-4.5-air',
      messages: [{ role: 'user', content: 'Hello' }],
      stream: false
    })
  });
  
  const response = await freshWorker.fetch(request, emptyPoolEnv);
  console.log('✓ Empty token pool status:', response.status);
  return response.status === 503;
}

async function testContentCleaning() {
  console.log('\n🧹 Testing content cleaning with upstream response sample...');

  try {
    // Import ProxyHandler directly
    const { ProxyHandler } = await import('./index.js');
    const fs = await import('fs/promises');

    // Create a mock handler to test cleaning functions
    const mockTokenManager = {
      getCommonHeaders: () => ({})
    };
    const mockRequestHandler = {
      tokenManager: mockTokenManager
    };
    const mockSettings = {};
    const handler = new ProxyHandler(mockRequestHandler, mockSettings);

    // Read upstream response sample
    const fixtureContent = await fs.readFile('./fixtures/upstream_response.txt', 'utf-8');

    // Parse SSE stream to extract thinking and answer content
    const lines = fixtureContent.split('\n');
    let thinkingParts = [];
    let answerParts = [];
    let currentPhase = null;

    for (const line of lines) {
      if (!line.trim().startsWith('data: ')) continue;

      const payloadStr = line.substring(6).trim();
      if (payloadStr === '[DONE]' || !payloadStr) continue;

      try {
        const data = JSON.parse(payloadStr).data || {};
        const phase = data.phase;
        const content = data.delta_content || data.edit_content;

        if (phase) currentPhase = phase;
        if (!content) continue;

        // Check for phase transition
        const match = content.match(/(.*<\/details>)(.*)/s);
        if (match) {
          const [, thinkPart, answerPart] = match;
          if (thinkPart) thinkingParts.push(thinkPart);
          if (answerPart) answerParts.push(answerPart);
          currentPhase = 'answer';
        } else {
          if (currentPhase === 'thinking') {
            thinkingParts.push(content);
          } else if (currentPhase === 'answer') {
            answerParts.push(content);
          }
        }
      } catch (e) {
        continue;
      }
    }

    // Test thinking content cleaning
    const rawThinking = thinkingParts.join('');
    const cleanedThinking = handler.cleanThinkingContent(rawThinking);

    // Test answer content cleaning
    const rawAnswer = answerParts.join('');
    const cleanedAnswer = handler.cleanAnswerContent(rawAnswer);

    // Read expected output
    const goldenTruthContent = await fs.readFile('./fixtures/expected_output.txt', 'utf-8');

    // Parse golden truth
    const thinkingMatch = goldenTruthContent.match(/CLEANED THINKING CONTENT:\n={80}\n(.*?)\n={80}/s);
    const answerMatch = goldenTruthContent.match(/CLEANED ANSWER CONTENT:\n={80}\n(.*?)$/s);

    const expectedThinking = thinkingMatch ? thinkingMatch[1].trim() : '';
    const expectedAnswer = answerMatch ? answerMatch[1].trim() : '';

    console.log('✓ Raw thinking length:', rawThinking.length);
    console.log('✓ Cleaned thinking length:', cleanedThinking.length);
    console.log('✓ Expected thinking length:', expectedThinking.length);
    console.log('✓ Raw answer length:', rawAnswer.length);
    console.log('✓ Cleaned answer length:', cleanedAnswer.length);
    console.log('✓ Expected answer length:', expectedAnswer.length);

    // Compare with golden truth
    const thinkingMatches = cleanedThinking.trim() === expectedThinking;
    const answerMatches = cleanedAnswer.trim() === expectedAnswer;

    console.log(`✓ Thinking matches golden truth: ${thinkingMatches ? '✅' : '❌'}`);
    console.log(`✓ Answer matches golden truth: ${answerMatches ? '✅' : '❌'}`);

    if (!thinkingMatches) {
      console.log('\n⚠️  Thinking content mismatch!');
      console.log('Expected (first 200 chars):', expectedThinking.substring(0, 200));
      console.log('Got (first 200 chars):', cleanedThinking.substring(0, 200));
    }

    if (!answerMatches) {
      console.log('\n⚠️  Answer content mismatch!');
      console.log('Expected:', expectedAnswer);
      console.log('Got:', cleanedAnswer);
    }

    // Verify content is properly cleaned
    const hasDetailsTag = cleanedThinking.includes('<details');
    const hasSummaryTag = cleanedThinking.includes('<summary');
    const hasGtPrefix = cleanedThinking.includes('\n> ');
    const answerClean = !cleanedAnswer.includes('<glm_block');

    console.log(`✓ No <details> tags: ${!hasDetailsTag}`);
    console.log(`✓ No <summary> tags: ${!hasSummaryTag}`);
    console.log(`✓ No "> " prefixes: ${!hasGtPrefix}`);
    console.log(`✓ No <glm_block> tags: ${answerClean}`);

    // Final verification
    const success = thinkingMatches && answerMatches &&
                   !hasDetailsTag && !hasSummaryTag && !hasGtPrefix && answerClean;

    console.log(`\n✓ Content cleaning test: ${success ? '✅ PASS' : '❌ FAIL'}`);
    return success;

  } catch (error) {
    console.log('✗ Content cleaning test failed:', error.message);
    console.log(error.stack);
    return false;
  }
}

async function runTests() {
  console.log('🧪 Running Z.AI Proxy Tests...\n');

  try {
    // Basic endpoint tests
    const healthOk = await testHealthEndpoint();
    const modelsOk = await testModelsEndpoint();

    // Content cleaning test
    const cleaningOk = await testContentCleaning();
    
    // Error handling tests
    const invalidAuthOk = await testInvalidAuth();
    const invalidModelOk = await testInvalidModel();
    const emptyTokenPoolOk = await testEmptyTokenPool();
    
    let chatOk = false;
    let streamingOk = false;
    
    // Test chat functionality with token pool
    console.log('\n🔄 Testing chat completion (this may take a moment)...');
    try {
      chatOk = await testChatCompletion();
    } catch (error) {
      if (error.message.includes('Token pool required')) {
        console.log('⚠️  Chat test skipped - token pool failed');
        chatOk = null;
      } else {
        console.log('⚠️  Chat test failed unexpectedly:', error.message);
        chatOk = false;
      }
    }
    
    console.log('\n🌊 Testing streaming chat (this may take a moment)...');
    try {
      streamingOk = await testStreamingChat();
    } catch (error) {
      if (error.message.includes('Token pool required')) {
        console.log('⚠️  Streaming test skipped - token pool failed');
        streamingOk = null;
      } else {
        console.log('⚠️  Streaming test failed unexpectedly:', error.message);
        streamingOk = false;
      }
    }
    
    console.log('\n📊 Test Results:');
    console.log(`Health endpoint: ${healthOk ? '✅ PASS' : '❌ FAIL'}`);
    console.log(`Models endpoint: ${modelsOk ? '✅ PASS' : '❌ FAIL'}`);
    console.log(`Content cleaning: ${cleaningOk ? '✅ PASS' : '❌ FAIL'}`);
    console.log(`Invalid auth handling: ${invalidAuthOk ? '✅ PASS' : '❌ FAIL'}`);
    console.log(`Invalid model handling: ${invalidModelOk ? '✅ PASS' : '❌ FAIL'}`);
    console.log(`Empty token pool handling: ${emptyTokenPoolOk ? '✅ PASS' : '❌ FAIL'}`);

    if (chatOk === null) {
      console.log(`Chat completion: ⏭️  SKIPPED`);
    } else {
      console.log(`Chat completion: ${chatOk ? '✅ PASS' : '❌ FAIL'}`);
    }

    if (streamingOk === null) {
      console.log(`Streaming chat: ⏭️  SKIPPED`);
    } else {
      console.log(`Streaming chat: ${streamingOk ? '✅ PASS' : '❌ FAIL'}`);
    }

    // Count only non-null results
    const coreTests = [healthOk, modelsOk, cleaningOk, invalidAuthOk, invalidModelOk, emptyTokenPoolOk];
    const chatTests = [chatOk, streamingOk].filter(result => result !== null);

    const totalTests = coreTests.length + chatTests.length;
    const passedTests = [...coreTests, ...chatTests].filter(Boolean).length;
    const skippedTests = [chatOk, streamingOk].filter(result => result === null).length;
    
    console.log(`\n🎯 Results: ${passedTests}/${totalTests} tests passed${skippedTests > 0 ? `, ${skippedTests} skipped` : ''}`);
    console.log('🎉 Tests completed!');
  } catch (error) {
    console.error('❌ Test failed:', error.message);
  }
}

runTests();