homelab-frontend/internal/proxy/bodydispatch_test.go

package proxy

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

	"github.com/Riotpiaole/homelab-frontend/internal/config"
)

// TestBodyBasedDispatch verifies that /v1/chat/completions routes based on model field.
func TestBodyBasedDispatch(t *testing.T) {
	// Create two separate upstreams to verify routing
	reasoningCalled := false
	ornithCalled := false

	reasoningServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		reasoningCalled = true
		if r.URL.Path != "/v1/chat/completions" {
			t.Errorf("expected path /v1/chat/completions, got %s", r.URL.Path)
		}
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusOK)
		fmt.Fprint(w, `{"choices":[{"message":{"content":"response from reasoning"}}]}`)
	}))
	defer reasoningServer.Close()

	ornithServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		ornithCalled = true
		if r.URL.Path != "/v1/chat/completions" {
			t.Errorf("expected path /v1/chat/completions, got %s", r.URL.Path)
		}
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusOK)
		fmt.Fprint(w, `{"choices":[{"message":{"content":"response from ornith"}}]}`)
	}))
	defer ornithServer.Close()

	reasoningAddr := strings.TrimPrefix(reasoningServer.URL, "http://")
	ornithAddr := strings.TrimPrefix(ornithServer.URL, "http://")

	cfg := &config.Config{
		Models: map[string]*config.ModelUpstream{
			"reasoning": {
				Name:    "reasoning",
				Address: reasoningAddr,
				Path:    "/v1/chat/completions",
			},
			"ornith:35b": {
				Name:    "ornith:35b",
				Address: ornithAddr,
				Path:    "/v1/chat/completions",
			},
			"qwen2.5:3b-instruct": {
				Name:    "qwen2.5:3b-instruct",
				Address: ornithAddr,
				Path:    "/v1/chat/completions",
			},
		},
		Routes: make(map[string]*config.Route),
	}

	handler := New(cfg)
	defer handler.Close()

	server := httptest.NewServer(handler)
	defer server.Close()

	// Test 1: Route to reasoning upstream
	reasoningCalled = false
	ornithCalled = false
	requestBody := `{"model":"reasoning","messages":[{"role":"user","content":"hi"}]}`
	resp, err := http.Post(server.URL+"/v1/chat/completions", "application/json", strings.NewReader(requestBody))
	if err != nil {
		t.Fatalf("request failed: %v", err)
	}
	resp.Body.Close()

	if !reasoningCalled {
		t.Errorf("expected reasoning upstream to be called")
	}
	if ornithCalled {
		t.Errorf("expected ornith upstream NOT to be called")
	}

	// Test 2: Route to ornith upstream (both models point there)
	reasoningCalled = false
	ornithCalled = false
	requestBody = `{"model":"ornith:35b","messages":[{"role":"user","content":"hi"}]}`
	resp, _ = http.Post(server.URL+"/v1/chat/completions", "application/json", strings.NewReader(requestBody))
	resp.Body.Close()

	if reasoningCalled {
		t.Errorf("expected reasoning upstream NOT to be called")
	}
	if !ornithCalled {
		t.Errorf("expected ornith upstream to be called")
	}

	// Test 3: qwen2.5 also routes to ornith
	reasoningCalled = false
	ornithCalled = false
	requestBody = `{"model":"qwen2.5:3b-instruct","messages":[{"role":"user","content":"hi"}]}`
	resp, _ = http.Post(server.URL+"/v1/chat/completions", "application/json", strings.NewReader(requestBody))
	resp.Body.Close()

	if reasoningCalled {
		t.Errorf("expected reasoning upstream NOT to be called")
	}
	if !ornithCalled {
		t.Errorf("expected ornith upstream to be called")
	}
}

// TestBodyPreservedUnmodified verifies that the request body is forwarded unmodified.
func TestBodyPreservedUnmodified(t *testing.T) {
	var receivedBody []byte

	upstreamServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var err error
		receivedBody, err = io.ReadAll(r.Body)
		if err != nil {
			t.Errorf("failed to read body: %v", err)
		}
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusOK)
		fmt.Fprint(w, `{}`)
	}))
	defer upstreamServer.Close()

	upstreamAddr := strings.TrimPrefix(upstreamServer.URL, "http://")

	cfg := &config.Config{
		Models: map[string]*config.ModelUpstream{
			"reasoning": {
				Name:    "reasoning",
				Address: upstreamAddr,
			},
		},
		Routes: make(map[string]*config.Route),
	}

	handler := New(cfg)
	defer handler.Close()

	server := httptest.NewServer(handler)
	defer server.Close()

	// Send a request with specific body content
	originalBody := `{"model":"reasoning","stream":true,"messages":[{"role":"user","content":"hello world"}],"temperature":0.7}`
	resp, err := http.Post(server.URL+"/v1/chat/completions", "application/json", strings.NewReader(originalBody))
	if err != nil {
		t.Fatalf("request failed: %v", err)
	}
	resp.Body.Close()

	if string(receivedBody) != originalBody {
		t.Errorf("body modified: expected %s, got %s", originalBody, string(receivedBody))
	}
}

// TestStreamingUnbuffered verifies that streaming responses are unbuffered.
func TestStreamingUnbuffered(t *testing.T) {
	upstreamServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "text/event-stream")
		w.Header().Set("Cache-Control", "no-cache")
		w.WriteHeader(http.StatusOK)

		rc := http.NewResponseController(w)
		for i := 0; i < 3; i++ {
			fmt.Fprintf(w, "data: {\"choices\":[{\"delta\":{\"content\":\"token%d\"}}]}\n\n", i)
			rc.Flush()
			time.Sleep(50 * time.Millisecond)
		}
		fmt.Fprint(w, "data: [DONE]\n\n")
		rc.Flush()
	}))
	defer upstreamServer.Close()

	upstreamAddr := strings.TrimPrefix(upstreamServer.URL, "http://")

	cfg := &config.Config{
		Models: map[string]*config.ModelUpstream{
			"reasoning": {
				Name:    "reasoning",
				Address: upstreamAddr,
			},
		},
		Routes: make(map[string]*config.Route),
	}

	handler := New(cfg)
	defer handler.Close()

	server := httptest.NewServer(handler)
	defer server.Close()

	requestBody := `{"model":"reasoning","stream":true,"messages":[]}`
	resp, err := http.Post(server.URL+"/v1/chat/completions", "application/json", strings.NewReader(requestBody))
	if err != nil {
		t.Fatalf("request failed: %v", err)
	}
	defer resp.Body.Close()

	if resp.Header.Get("Content-Type") != "text/event-stream" {
		t.Errorf("expected Content-Type: text/event-stream, got %s", resp.Header.Get("Content-Type"))
	}

	// Read response and verify streaming
	startTime := time.Now()
	respBody, _ := io.ReadAll(resp.Body)
	duration := time.Since(startTime)

	respStr := string(respBody)
	if !strings.Contains(respStr, "token0") || !strings.Contains(respStr, "token1") || !strings.Contains(respStr, "token2") {
		t.Errorf("expected all tokens in response, got: %s", respStr)
	}

	// With 50ms gaps and 3 tokens, we should take at least 100ms
	// If completely buffered, would be much faster
	if duration < 100*time.Millisecond {
		t.Logf("response arrived very quickly (%.0fms) - may indicate buffering", duration.Seconds()*1000)
	}
}

// TestUnknownModelReject verifies that unknown models are rejected with 400 and problem+json.
func TestUnknownModelReject(t *testing.T) {
	upstreamServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
	}))
	defer upstreamServer.Close()

	upstreamAddr := strings.TrimPrefix(upstreamServer.URL, "http://")

	cfg := &config.Config{
		Models: map[string]*config.ModelUpstream{
			"reasoning": {
				Name:    "reasoning",
				Address: upstreamAddr,
			},
		},
		Routes: make(map[string]*config.Route),
	}

	handler := New(cfg)
	defer handler.Close()

	server := httptest.NewServer(handler)
	defer server.Close()

	requestBody := `{"model":"unknown-model","messages":[]}`
	resp, err := http.Post(server.URL+"/v1/chat/completions", "application/json", strings.NewReader(requestBody))
	if err != nil {
		t.Fatalf("request failed: %v", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusBadRequest {
		t.Errorf("expected 400 for unknown model, got %d", resp.StatusCode)
	}

	// Verify problem+json content type
	ct := resp.Header.Get("Content-Type")
	if !strings.Contains(ct, "application/problem+json") {
		t.Errorf("expected content-type application/problem+json, got %s", ct)
	}

	// Verify response is valid JSON
	var prob map[string]interface{}
	if err := json.NewDecoder(resp.Body).Decode(&prob); err != nil {
		t.Errorf("response is not valid JSON: %v", err)
	}

	// Verify valid_models is included
	if prob["valid_models"] == nil {
		t.Errorf("expected valid_models in problem detail")
	}
}

// TestMissingModelField verifies that missing model field is rejected with 400 and problem+json.
func TestMissingModelField(t *testing.T) {
	upstreamServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
	}))
	defer upstreamServer.Close()

	upstreamAddr := strings.TrimPrefix(upstreamServer.URL, "http://")

	cfg := &config.Config{
		Models: map[string]*config.ModelUpstream{
			"reasoning": {
				Name:    "reasoning",
				Address: upstreamAddr,
			},
		},
		Routes: make(map[string]*config.Route),
	}

	handler := New(cfg)
	defer handler.Close()

	server := httptest.NewServer(handler)
	defer server.Close()

	requestBody := `{"messages":[]}`
	resp, err := http.Post(server.URL+"/v1/chat/completions", "application/json", strings.NewReader(requestBody))
	if err != nil {
		t.Fatalf("request failed: %v", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusBadRequest {
		t.Errorf("expected 400 for missing model, got %d", resp.StatusCode)
	}

	// Verify problem+json content type
	ct := resp.Header.Get("Content-Type")
	if !strings.Contains(ct, "application/problem+json") {
		t.Errorf("expected content-type application/problem+json, got %s", ct)
	}

	// Verify response is valid JSON
	var prob map[string]interface{}
	if err := json.NewDecoder(resp.Body).Decode(&prob); err != nil {
		t.Errorf("response is not valid JSON: %v", err)
	}
}

// TestBodySize verifies that body size cap is enforced for dispatched requests.
func TestBodySizeCappedDispatch(t *testing.T) {
	upstreamServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
	}))
	defer upstreamServer.Close()

	upstreamAddr := strings.TrimPrefix(upstreamServer.URL, "http://")

	cfg := &config.Config{
		Models: map[string]*config.ModelUpstream{
			"reasoning": {
				Name:    "reasoning",
				Address: upstreamAddr,
			},
		},
		Routes: make(map[string]*config.Route),
	}

	handler := New(cfg)
	defer handler.Close()

	server := httptest.NewServer(handler)
	defer server.Close()

	// Create a JSON payload
	payload := map[string]interface{}{
		"model": "reasoning",
		"messages": []map[string]string{
			{
				"role":    "user",
				"content": strings.Repeat("x", 100),
			},
		},
	}

	body, _ := json.Marshal(payload)
	req, _ := http.NewRequest("POST", server.URL+"/v1/chat/completions", bytes.NewReader(body))
	req.Header.Set("Content-Type", "application/json")
	req.ContentLength = int64(len(body))

	resp, _ := http.DefaultClient.Do(req)
	resp.Body.Close()

	// Request should be accepted (body size is reasonable)
	if resp.StatusCode != http.StatusOK {
		t.Errorf("expected 200 for reasonable body, got %d", resp.StatusCode)
	}
}