ModelRouter/internal/provider/provider_test.go

package provider

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"net/http/httptest"
	"sync"
	"testing"
	"time"

	"llmsproxy/internal/config"
	"llmsproxy/internal/lua"
	"llmsproxy/internal/types"
)

func newTestProvider(t *testing.T, src config.Source) *Provider {
	t.Helper()
	vm := lua.NewVM(t.TempDir())
	if err := vm.Start(); err != nil {
		t.Fatalf("vm: %v", err)
	}
	t.Cleanup(vm.Stop)
	return New(src, vm)
}

func src(name, url, adapter string, models ...string) config.Source {
	s := config.Source{Name: name, BaseURL: url, Adapter: adapter, MaxConcurrent: 4}
	for _, m := range models {
		s.Models = append(s.Models, config.Model{ID: m, Priority: 0})
	}
	return s
}

func TestProviderChat(t *testing.T) {
	up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var body map[string]interface{}
		_ = json.NewDecoder(r.Body).Decode(&body)
		if body["model"] != "mock-model" {
			t.Errorf("model = %v", body["model"])
		}
		fmt.Fprint(w, `{"choices":[{"message":{"content":"hi"},"finish_reason":"stop"}]}`)
	}))
	defer up.Close()
	p := newTestProvider(t, src("mock", up.URL, "openai", "mock-model"))
	resp, err := p.Chat(context.Background(), &types.ChatRequest{
		Model:    "mock-model",
		Messages: []types.ChatMessage{{Role: "user", Content: types.StringContent("hello")}},
	})
	if err != nil {
		t.Fatalf("chat: %v", err)
	}
	if resp.Content != "hi" {
		t.Fatalf("content = %q", resp.Content)
	}
}

func TestProviderChatStream(t *testing.T) {
	up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "text/event-stream")
		fmt.Fprint(w, "data: {\"choices\":[{\"delta\":{\"content\":\"a\"}}]}\n\n")
		fmt.Fprint(w, "data: {\"choices\":[{\"delta\":{\"content\":\"b\"}}]}\n\n")
		fmt.Fprint(w, "data: {\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}]}\n\n")
		fmt.Fprint(w, "data: [DONE]\n\n")
	}))
	defer up.Close()
	p := newTestProvider(t, src("mock", up.URL, "openai", "m"))
	ch, err := p.ChatStream(context.Background(), &types.ChatRequest{
		Model:    "m",
		Messages: []types.ChatMessage{{Role: "user", Content: types.StringContent("x")}},
	})
	if err != nil {
		t.Fatalf("stream: %v", err)
	}
	var text string
	var done bool
	for ck := range ch {
		text += ck.Content
		if ck.Done {
			done = true
		}
	}
	if text != "ab" {
		t.Fatalf("text = %q", text)
	}
	if !done {
		t.Fatal("expected done")
	}
}

func TestProviderImage(t *testing.T) {
	up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprint(w, `{"created":123,"data":[{"b64_json":"QUJD"}]}`)
	}))
	defer up.Close()
	p := newTestProvider(t, src("img", up.URL, "openai", "flux-1"))
	resp, err := p.Image(context.Background(), &types.ImageGenRequest{Model: "flux-1", Prompt: "cat"})
	if err != nil {
		t.Fatalf("image: %v", err)
	}
	if len(resp.ImageData) != 1 || resp.ImageData[0].B64JSON != "QUJD" {
		t.Fatalf("image data = %+v", resp.ImageData)
	}
}

func TestProviderBackoff(t *testing.T) {
	up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(500)
		fmt.Fprint(w, "boom")
	}))
	defer up.Close()
	p := newTestProvider(t, src("mock", up.URL, "openai", "m"))
	_, err := p.Chat(context.Background(), &types.ChatRequest{
		Model:    "m",
		Messages: []types.ChatMessage{{Role: "user", Content: types.StringContent("x")}},
	})
	if err == nil {
		t.Fatal("expected error")
	}
	if p.Available() {
		t.Fatal("expected provider to be in backoff")
	}
	// 401 -> permanent
	up2 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(401)
	}))
	defer up2.Close()
	p2 := newTestProvider(t, src("mock2", up2.URL, "openai", "m2"))
	p2.Chat(context.Background(), &types.ChatRequest{Model: "m2", Messages: []types.ChatMessage{{Role: "user", Content: types.StringContent("x")}}})
	if p2.Available() {
		t.Fatal("expected permanent unavailability on 401")
	}
}

func TestProviderConcurrencyCap(t *testing.T) {
	release := make(chan struct{})
	started := make(chan struct{}, 100)
	up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		started <- struct{}{}
		<-release
		fmt.Fprint(w, `{"choices":[{"message":{"content":"ok"}}]}`)
	}))
	defer up.Close()
	// cap 2
	p := newTestProvider(t, src("mock", up.URL, "openai", "m"))
	p.cfg.MaxConcurrent = 2
	p.sem = make(chan struct{}, 2)

	var wg sync.WaitGroup
	for i := 0; i < 6; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			p.Chat(context.Background(), &types.ChatRequest{Model: "m", Messages: []types.ChatMessage{{Role: "user", Content: types.StringContent("x")}}})
		}()
	}
	// wait until 2 requests started
	deadline := time.Now().Add(2 * time.Second)
	for len(started) < 2 {
		if time.Now().After(deadline) {
			t.Fatal("timeout waiting for first two")
		}
		time.Sleep(5 * time.Millisecond)
	}
	time.Sleep(100 * time.Millisecond)
	if len(started) > 2 {
		t.Fatalf("more than 2 concurrent: %d", len(started))
	}
	close(release)
	wg.Wait()
}