claudish/tests/zai-handler.test.ts

import { describe, test, expect, beforeAll, afterAll } from "bun:test";
import { createProxyServer } from "../src/proxy-server.js";

const ZAI_API_KEY = process.env.ZAI_API_KEY;
const TEST_PORT = 3456;

describe("ZaiHandler", () => {
  let proxy: { port: number; url: string; shutdown: () => Promise<void> } | null = null;

  beforeAll(async () => {
    if (!ZAI_API_KEY) {
      console.log("Skipping Z.ai tests - ZAI_API_KEY not set");
      return;
    }
    proxy = await createProxyServer(TEST_PORT, undefined, "z-ai/glm-4.6");
  });

  afterAll(async () => {
    if (proxy) {
      await proxy.shutdown();
    }
  });

  test("should detect z-ai model prefix", () => {
    const models = [
      "z-ai/glm-4.6",
      "z-ai/glm-4",
      "z-ai/glm-4-flash",
    ];

    for (const model of models) {
      expect(model.startsWith("z-ai/")).toBe(true);
    }
  });

  test("should convert z-ai/ prefix to model name", () => {
    const testCases = [
      { input: "z-ai/glm-4.6", expected: "glm-4.6" },
      { input: "z-ai/glm-4", expected: "glm-4" },
      { input: "z-ai/glm-4-flash", expected: "glm-4-flash" },
    ];

    for (const { input, expected } of testCases) {
      const result = input.startsWith("z-ai/") ? input.slice(5) : input;
      expect(result).toBe(expected);
    }
  });

  test("should make request to Z.ai API", async () => {
    if (!ZAI_API_KEY || !proxy) {
      console.log("Skipping - ZAI_API_KEY not set");
      return;
    }

    const response = await fetch(`${proxy.url}/v1/messages`, {
      method: "POST",
      headers: {
        "Content-Type": "application/json",
      },
      body: JSON.stringify({
        model: "z-ai/glm-4.6",
        max_tokens: 100,
        messages: [
          {
            role: "user",
            content: "Say 'Hello from Z.ai!' and nothing else."
          }
        ]
      })
    });

    expect(response.ok).toBe(true);
    expect(response.headers.get("content-type")).toContain("text/event-stream");

    // Read streaming response
    const reader = response.body!.getReader();
    const decoder = new TextDecoder();
    let fullText = "";
    let hasMessageStart = false;
    let hasMessageStop = false;

    while (true) {
      const { done, value } = await reader.read();
      if (done) break;

      const chunk = decoder.decode(value, { stream: true });
      fullText += chunk;

      if (chunk.includes("message_start")) hasMessageStart = true;
      if (chunk.includes("message_stop")) hasMessageStop = true;
    }

    expect(hasMessageStart).toBe(true);
    expect(hasMessageStop).toBe(true);
    console.log("Z.ai response received successfully");
  }, 30000);
});

// Standalone test runner
if (import.meta.main) {
  const ZAI_KEY = process.env.ZAI_API_KEY;

  if (!ZAI_KEY) {
    console.error("Error: ZAI_API_KEY environment variable is not set");
    console.log("\nUsage: ZAI_API_KEY=your_key bun tests/zai-handler.test.ts");
    process.exit(1);
  }

  console.log("Running Z.ai handler test...");
  console.log(`API Key: ${ZAI_KEY.slice(0, 8)}...${ZAI_KEY.slice(-4)}`);

  // Simple direct API test
  const testDirectApi = async () => {
    console.log("\n1. Testing direct Z.ai API call...");

    const response = await fetch("https://api.z.ai/api/coding/paas/v4/chat/completions", {
      method: "POST",
      headers: {
        "Content-Type": "application/json",
        "Authorization": `Bearer ${ZAI_KEY}`,
      },
      body: JSON.stringify({
        model: "glm-4.6",
        max_tokens: 100,
        stream: true,
        messages: [
          { role: "user", content: "Say 'Hello from GLM!' and nothing else." }
        ]
      })
    });

    if (!response.ok) {
      const error = await response.text();
      console.error(`API Error (${response.status}):`, error);
      return false;
    }

    console.log("Response status:", response.status);
    console.log("Content-Type:", response.headers.get("content-type"));

    const reader = response.body!.getReader();
    const decoder = new TextDecoder();
    let output = "";

    while (true) {
      const { done, value } = await reader.read();
      if (done) break;
      const chunk = decoder.decode(value, { stream: true });
      output += chunk;
    }

    console.log("\nRaw response sample:", output.slice(0, 500));
    return true;
  };

  // Test through proxy
  const testThroughProxy = async () => {
    console.log("\n2. Testing through Claudish proxy...");

    const proxy = await createProxyServer(3457, undefined, "z-ai/glm-4.6");
    console.log(`Proxy started at ${proxy.url}`);

    try {
      const response = await fetch(`${proxy.url}/v1/messages`, {
        method: "POST",
        headers: { "Content-Type": "application/json" },
        body: JSON.stringify({
          model: "z-ai/glm-4.6",
          max_tokens: 100,
          messages: [{ role: "user", content: "Say 'Hello through proxy!' and nothing else." }]
        })
      });

      if (!response.ok) {
        const error = await response.text();
        console.error(`Proxy Error (${response.status}):`, error);
        return false;
      }

      const reader = response.body!.getReader();
      const decoder = new TextDecoder();
      let output = "";
      let textContent = "";

      while (true) {
        const { done, value } = await reader.read();
        if (done) break;
        const chunk = decoder.decode(value, { stream: true });
        output += chunk;

        // Extract text deltas
        const matches = chunk.matchAll(/text_delta.*?"text":\s*"([^"]*)"/g);
        for (const match of matches) {
          textContent += match[1];
        }
      }

      console.log("\nExtracted text content:", textContent || "(no text found)");
      console.log("Full SSE events received:", output.includes("message_start") && output.includes("message_stop") ? "Yes" : "No");
      return true;
    } finally {
      await proxy.shutdown();
      console.log("Proxy shutdown complete");
    }
  };

  // Run tests
  (async () => {
    const directOk = await testDirectApi();
    if (directOk) {
      await testThroughProxy();
    }
    console.log("\n✓ Test complete");
  })();
}