qorona/qorona_test.go

// Unit and integration tests for qorona
// Created by Luxferre in 2026, released into the public domain

package main

import (
  "encoding/json"
  "net/http"
  "net/http/httptest"
  "strings"
  "testing"
  "time"
)

func TestChatMessageGetContentString(t *testing.T) {
  msg1 := ChatMessage{Role: "user", Content: "Hello world"}
  if msg1.GetContentString() != "Hello world" {
    t.Fatalf("expected 'Hello world', got %q", msg1.GetContentString())
  }

  msg2 := ChatMessage{
    Role: "user",
    Content: []interface{}{
      map[string]interface{}{"type": "text", "text": "Part 1 "},
      map[string]interface{}{"type": "text", "text": "Part 2"},
    },
  }
  if msg2.GetContentString() != "Part 1 Part 2" {
    t.Fatalf("expected 'Part 1 Part 2', got %q", msg2.GetContentString())
  }
}

func TestExtractThinkingContent(t *testing.T) {
  raw := "<think>\nAnalyzing the user's request...\n</think>\nHere is the answer."
  thinking, clean := ExtractThinkingContent(raw)
  if thinking != "Analyzing the user's request..." {
    t.Fatalf("unexpected thinking extraction: %q", thinking)
  }
  if clean != "Here is the answer." {
    t.Fatalf("unexpected clean text: %q", clean)
  }
}

func TestDetectToolCalls(t *testing.T) {
  xmlInput := "Let me check the weather.\n<tool_call>{\"name\": \"get_weather\", \"arguments\": {\"location\": \"Tokyo\"}}</tool_call>"
  calls, rem := DetectToolCalls(xmlInput)
  if len(calls) != 1 {
    t.Fatalf("expected 1 tool call, got %d", len(calls))
  }
  if calls[0].Function.Name != "get_weather" {
    t.Fatalf("expected function name 'get_weather', got %q", calls[0].Function.Name)
  }
  if !strings.Contains(calls[0].Function.Arguments, "Tokyo") {
    t.Fatalf("expected argument with Tokyo, got %q", calls[0].Function.Arguments)
  }
  if strings.TrimSpace(rem) != "Let me check the weather." {
    t.Fatalf("unexpected remaining text: %q", rem)
  }
}

func TestPrepareConversation(t *testing.T) {
  req := ChatCompletionRequest{
    Model: "qwen-3.8-max",
    Messages: []ChatMessage{
      {Role: "system", Content: "You are a helpful assistant."},
      {Role: "user", Content: "Tell me a joke."},
      {Role: "assistant", Content: "Why did the chicken cross the road?"},
      {Role: "user", Content: "Why?"},
    },
    ReasoningEffort: "medium",
  }

  history, question, thinkingMode := PrepareConversation(req)
  if thinkingMode != "true" {
    t.Fatalf("expected thinkingMode 'true', got %q", thinkingMode)
  }
  if len(history) != 2 {
    t.Fatalf("expected 2 history items, got %d", len(history))
  }
  if question != "Why?" {
    t.Fatalf("expected question 'Why?', got %q", question)
  }
  if !strings.Contains(history[0]["content"].(string), "You are a helpful assistant.") {
    t.Fatalf("expected system prompt inside first turn, got %v", history[0]["content"])
  }
}

func TestModelsHandler(t *testing.T) {
  gw := NewQoronaGateway("https://mock.hf.space", "/dummy/path.jpg", "qwen-3.8-max", 5*time.Second)

  req := httptest.NewRequest("GET", "/v1/models", nil)
  w := httptest.NewRecorder()

  gw.HandleModels(w, req)

  if w.Code != http.StatusOK {
    t.Fatalf("expected 200 OK, got %d", w.Code)
  }

  var res ModelsResponse
  if err := json.NewDecoder(w.Body).Decode(&res); err != nil {
    t.Fatalf("failed to decode response: %v", err)
  }

  if len(res.Data) == 0 {
    t.Fatalf("expected at least 1 model in response")
  }

  found := false
  for _, m := range res.Data {
    if m.ID == "qorona" || m.ID == "qwen-3.8-max" {
      found = true
      break
    }
  }
  if !found {
    t.Fatalf("model not found in models list")
  }
}

func TestBuildUsage(t *testing.T) {
  u := BuildUsage(100, 50, 80, 20)
  if u.PromptTokens != 100 {
    t.Fatalf("expected prompt_tokens 100, got %d", u.PromptTokens)
  }
  if u.CompletionTokens != 50 {
    t.Fatalf("expected completion_tokens 50, got %d", u.CompletionTokens)
  }
  if u.TotalTokens != 150 {
    t.Fatalf("expected total_tokens 150, got %d", u.TotalTokens)
  }
  if u.PromptTokensDetails == nil || u.PromptTokensDetails.CachedTokens != 80 {
    t.Fatalf("expected cached_tokens 80, got %+v", u.PromptTokensDetails)
  }
  if u.CompletionTokensDetails == nil || u.CompletionTokensDetails.ReasoningTokens != 20 {
    t.Fatalf("expected reasoning_tokens 20, got %+v", u.CompletionTokensDetails)
  }

  b, err := json.Marshal(u)
  if err != nil {
    t.Fatalf("failed to marshal usage: %v", err)
  }
  if !strings.Contains(string(b), `"cached_tokens":80`) {
    t.Fatalf("expected json to contain cached_tokens: %s", string(b))
  }
}

func TestWriteCompletionResponseUsage(t *testing.T) {
  w := httptest.NewRecorder()
  out := FinalOutput{
    Content:          "Hello",
    ReasoningContent: "Thinking...",
    FinishReason:     "stop",
    PromptTokens:     900,
    CompletionTokens: 25,
    CachedTokens:     850,
    ReasoningTokens:  10,
  }

  WriteCompletionResponse(w, "chatcmpl-test", 1234567890, "qwen-3.8-max", out)
  if w.Code != http.StatusOK {
    t.Fatalf("expected 200 OK, got %d", w.Code)
  }

  var resp ChatCompletionResponse
  if err := json.NewDecoder(w.Body).Decode(&resp); err != nil {
    t.Fatalf("failed to decode response: %v", err)
  }

  if resp.Usage.PromptTokens != 900 {
    t.Fatalf("expected prompt_tokens 900, got %d", resp.Usage.PromptTokens)
  }
  if resp.Usage.PromptTokensDetails == nil || resp.Usage.PromptTokensDetails.CachedTokens != 850 {
    t.Fatalf("expected cached_tokens 850, got %+v", resp.Usage.PromptTokensDetails)
  }
}

func TestStreamerUsage(t *testing.T) {
  w := httptest.NewRecorder()
  streamer := NewStreamer(w, nil, "chatcmpl-stream-test", 1234567890, "qwen-3.8-max")
  usage := BuildUsage(900, 30, 850, 0)

  streamer.EmitUsage(usage)
  streamer.Done()

  body := w.Body.String()
  if !strings.Contains(body, `"cached_tokens":850`) {
    t.Fatalf("expected streamed chunk to contain cached_tokens: %s", body)
  }
  if !strings.Contains(body, `data: [DONE]`) {
    t.Fatalf("expected streamed body to contain [DONE]: %s", body)
  }
}

func TestResolveCachedTokens(t *testing.T) {
  // If upstream explicitly provides cached tokens (>0), preserve it
  if got := ResolveCachedTokens(500, 1000, "hello"); got != 500 {
    t.Fatalf("expected 500, got %d", got)
  }

  // If upstream returns 0, estimate based on promptTokens and current question
  if got := ResolveCachedTokens(0, 1397, "Say hello in 1 word"); got <= 0 || got >= 1397 {
    t.Fatalf("expected cached tokens between 1 and 1396, got %d", got)
  }

  // If promptTokens <= 0, return 0
  if got := ResolveCachedTokens(0, 0, "hello"); got != 0 {
    t.Fatalf("expected 0, got %d", got)
  }
}