descrybe/apps/api/cmd/mock-llm/main_test.go

package main

import (
	"bytes"
	"context"
	"encoding/json"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	"github.com/descrybe/descrybe-v2/apps/api/internal/processing"
)

func testServer(t *testing.T) *httptest.Server {
	t.Helper()
	s := &server{apiKey: "local-test", model: "mock-llm"}
	mux := http.NewServeMux()
	mux.HandleFunc("/healthz", s.handleHealth)
	mux.HandleFunc("/v1/models", s.handleModels)
	mux.HandleFunc("/v1/chat/completions", s.handleChatCompletions)
	mux.HandleFunc("/v1/embeddings", s.handleEmbeddings)
	srv := httptest.NewServer(mux)
	t.Cleanup(srv.Close)
	return srv
}

func TestMockLLM_healthAndModels(t *testing.T) {
	t.Parallel()
	srv := testServer(t)

	res, err := http.Get(srv.URL + "/healthz")
	if err != nil {
		t.Fatal(err)
	}
	defer res.Body.Close()
	if res.StatusCode != http.StatusOK {
		t.Fatalf("health status=%d", res.StatusCode)
	}

	req, err := http.NewRequest(http.MethodGet, srv.URL+"/v1/models", nil)
	if err != nil {
		t.Fatal(err)
	}
	req.Header.Set("Authorization", "Bearer local-test")
	res2, err := http.DefaultClient.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	defer res2.Body.Close()
	if res2.StatusCode != http.StatusOK {
		t.Fatalf("models status=%d", res2.StatusCode)
	}
	var body map[string]any
	if err := json.NewDecoder(res2.Body).Decode(&body); err != nil {
		t.Fatal(err)
	}
	data, _ := body["data"].([]any)
	if len(data) < 1 {
		t.Fatalf("models empty: %#v", body)
	}
}

func TestMockLLM_chatCompletionsEnhanceJSON(t *testing.T) {
	t.Parallel()
	srv := testServer(t)

	payload := map[string]any{
		"model": "mock-llm",
		"messages": []map[string]string{
			{"role": "system", "content": `Return JSON with "name" and "description" for titles and descriptions.`},
			{"role": "user", "content": "current name: Red Runner\ncurrent description: A fine shoe.\nattributes:"},
		},
		"temperature": 0.2,
		"max_tokens":  350,
	}
	raw, err := json.Marshal(payload)
	if err != nil {
		t.Fatal(err)
	}
	req, err := http.NewRequest(http.MethodPost, srv.URL+"/v1/chat/completions", bytes.NewReader(raw))
	if err != nil {
		t.Fatal(err)
	}
	req.Header.Set("Authorization", "Bearer local-test")
	req.Header.Set("Content-Type", "application/json")
	res, err := http.DefaultClient.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	defer res.Body.Close()
	body, err := io.ReadAll(res.Body)
	if err != nil {
		t.Fatal(err)
	}
	if res.StatusCode != http.StatusOK {
		t.Fatalf("status=%d body=%s", res.StatusCode, body)
	}
	var parsed struct {
		Model   string `json:"model"`
		Choices []struct {
			Message struct {
				Content string `json:"content"`
			} `json:"message"`
		} `json:"choices"`
	}
	if err := json.Unmarshal(body, &parsed); err != nil {
		t.Fatal(err)
	}
	if parsed.Model != "mock-llm" {
		t.Fatalf("model=%q", parsed.Model)
	}
	if len(parsed.Choices) < 1 {
		t.Fatal("no choices")
	}
	content := parsed.Choices[0].Message.Content
	var obj map[string]any
	if err := json.Unmarshal([]byte(content), &obj); err != nil {
		t.Fatalf("content not JSON: %q err=%v", content, err)
	}
	if name, _ := obj["name"].(string); name == "" {
		t.Fatalf("missing name in %#v", obj)
	}
	if desc, _ := obj["description"].(string); desc == "" {
		t.Fatalf("missing description in %#v", obj)
	}
}

func TestMockLLM_OpenAIClientRoundTrip(t *testing.T) {
	t.Parallel()
	srv := testServer(t)

	client := processing.NewOpenAIClient("local-test", srv.URL+"/v1", "mock-llm", 0, 1)
	if !client.Enabled() {
		t.Fatal("expected Enabled")
	}
	comp, err := client.Complete(context.Background(),
		`Return JSON with "name" and titles and descriptions.`,
		"current name: Mock Widget\ncurrent description: Tiny fixture.\nattributes:",
	)
	if err != nil {
		t.Fatal(err)
	}
	if !strings.Contains(comp.Text, "name") {
		t.Fatalf("unexpected text=%q", comp.Text)
	}
	if comp.TotalTokens < 1 {
		t.Fatalf("tokens=%d", comp.TotalTokens)
	}
}

func TestMockLLM_rejectsBadAuth(t *testing.T) {
	t.Parallel()
	srv := testServer(t)
	req, err := http.NewRequest(http.MethodGet, srv.URL+"/v1/models", nil)
	if err != nil {
		t.Fatal(err)
	}
	req.Header.Set("Authorization", "Bearer wrong")
	res, err := http.DefaultClient.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	defer res.Body.Close()
	if res.StatusCode != http.StatusUnauthorized {
		t.Fatalf("status=%d", res.StatusCode)
	}
}