vojo/apps/ai-bot/failover_test.go

package main

import (
	"bytes"
	"context"
	"errors"
	"log/slog"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"
)

// funcLLM is an LLMClient driven by a closure — for failover legs that need to
// block, inspect their context, or fail in a specific typed way.
type funcLLM struct {
	fn    func(ctx context.Context, req LLMRequest) (*LLMResponse, error)
	calls int
}

func (f *funcLLM) Complete(ctx context.Context, req LLMRequest) (*LLMResponse, error) {
	f.calls++
	return f.fn(ctx, req)
}

// newTestFailover builds the decorator without the background prober, so tests
// control the breaker state deterministically.
func newTestFailover(primary, fallback LLMClient, healthy bool) *failoverLLMClient {
	return &failoverLLMClient{
		primary:    primary,
		fallback:   fallback,
		log:        discardLog(),
		legTimeout: time.Second,
		healthy:    healthy,
		needOK:     1,
	}
}

// TestFailoverHealthyPrefersLocal: with a healthy breaker the local leg answers and
// the cloud is never touched; the response carries the local leg's Model untouched.
func TestFailoverHealthyPrefersLocal(t *testing.T) {
	local := &fakeLLM{text: "local answer", model: "qwen-local"}
	cloud := &fakeLLM{text: "cloud answer", model: "grok-x"}
	c := newTestFailover(local, cloud, true)

	resp, err := c.Complete(context.Background(), LLMRequest{Model: "grok-x"})
	if err != nil {
		t.Fatalf("Complete: %v", err)
	}
	if resp.Text != "local answer" || resp.Model != "qwen-local" {
		t.Fatalf("resp = (%q, model %q), want the local leg's answer", resp.Text, resp.Model)
	}
	if local.calls != 1 || cloud.calls != 0 {
		t.Fatalf("calls local=%d cloud=%d, want 1/0", local.calls, cloud.calls)
	}
	if !c.isHealthy() {
		t.Fatal("a successful local answer must not touch the breaker")
	}
}

// TestFailoverUnhealthyGoesStraightToCloud: an open breaker sends the request to the
// cloud without paying the local leg's timeout (the local client is never called).
func TestFailoverUnhealthyGoesStraightToCloud(t *testing.T) {
	local := &fakeLLM{text: "local"}
	cloud := &fakeLLM{text: "cloud answer", model: "grok-x"}
	c := newTestFailover(local, cloud, false)

	resp, err := c.Complete(context.Background(), LLMRequest{})
	if err != nil {
		t.Fatalf("Complete: %v", err)
	}
	if resp.Text != "cloud answer" {
		t.Fatalf("resp = %q, want the cloud answer", resp.Text)
	}
	if local.calls != 0 || cloud.calls != 1 {
		t.Fatalf("calls local=%d cloud=%d, want 0/1", local.calls, cloud.calls)
	}
}

// TestFailoverTransportErrorFallsBackAndTrips: a network-class local failure answers
// from the cloud within the SAME request and opens the breaker, so the next request
// skips the local leg entirely.
func TestFailoverTransportErrorFallsBackAndTrips(t *testing.T) {
	local := &fakeLLM{err: errors.New("dial tcp: connection refused")}
	cloud := &fakeLLM{text: "cloud answer", model: "grok-x"}
	c := newTestFailover(local, cloud, true)

	resp, err := c.Complete(context.Background(), LLMRequest{})
	if err != nil || resp.Text != "cloud answer" {
		t.Fatalf("resp=(%v,%v), want the cloud fallback", resp, err)
	}
	if c.isHealthy() {
		t.Fatal("a transport failure must trip the breaker")
	}
	if _, err := c.Complete(context.Background(), LLMRequest{}); err != nil {
		t.Fatalf("second Complete: %v", err)
	}
	if local.calls != 1 || cloud.calls != 2 {
		t.Fatalf("calls local=%d cloud=%d, want 1/2 (tripped breaker skips local)", local.calls, cloud.calls)
	}
}

// TestFailover5xxAnd429FallBack: retryable-class statuses (a dying tunnel's 502, a
// busy single-slot GPU's 429) fall back to the cloud like a transport error.
func TestFailover5xxAnd429FallBack(t *testing.T) {
	for _, status := range []int{500, 502, http.StatusTooManyRequests} {
		local := &fakeLLM{err: &httpStatusError{provider: "local", status: status, body: "boom"}}
		cloud := &fakeLLM{text: "cloud answer"}
		c := newTestFailover(local, cloud, true)
		resp, err := c.Complete(context.Background(), LLMRequest{})
		if err != nil || resp.Text != "cloud answer" {
			t.Fatalf("status %d: resp=(%v,%v), want the cloud fallback", status, resp, err)
		}
	}
}

// TestFailoverTerminal4xxFailsLoud: a local 4xx is a config/request error (a wrong
// LOCAL_LLM_MODEL tag) — masking it with a paid cloud answer would hide the
// misconfiguration, so it surfaces as the request's error and the cloud is not called.
func TestFailoverTerminal4xxFailsLoud(t *testing.T) {
	local := &fakeLLM{err: &httpStatusError{provider: "local", status: 404, body: "model not found"}}
	cloud := &fakeLLM{text: "cloud"}
	c := newTestFailover(local, cloud, true)

	if _, err := c.Complete(context.Background(), LLMRequest{}); err == nil {
		t.Fatal("want the local 4xx surfaced, got nil")
	}
	if cloud.calls != 0 {
		t.Fatalf("cloud called %d times on a local 4xx, want 0 (fail loud)", cloud.calls)
	}
	if !c.isHealthy() {
		t.Fatal("a 4xx is not an availability failure — the breaker must stay closed")
	}
}

// TestFailoverBudgetGoneNoCloudRetry: when the WHOLE request budget is exhausted
// (parent ctx done), the local error surfaces without burning a doomed cloud attempt
// and without tripping the breaker over our own cancellation.
func TestFailoverBudgetGoneNoCloudRetry(t *testing.T) {
	local := &funcLLM{fn: func(ctx context.Context, _ LLMRequest) (*LLMResponse, error) {
		<-ctx.Done()
		return nil, ctx.Err()
	}}
	cloud := &fakeLLM{text: "cloud"}
	c := newTestFailover(local, cloud, true)

	ctx, cancel := context.WithCancel(context.Background())
	cancel() // the shared RequestBudget is gone before the call
	if _, err := c.Complete(ctx, LLMRequest{}); err == nil {
		t.Fatal("want an error when the request budget is gone")
	}
	if cloud.calls != 0 {
		t.Fatalf("cloud called %d times after budget exhaustion, want 0", cloud.calls)
	}
	if !c.isHealthy() {
		t.Fatal("our own cancellation must not trip the breaker")
	}
}

// TestFailoverLocalTimeoutFallsBack: the local leg's own deadline (a half-alive
// tunnel that accepts but never answers) expires and the request is served by the
// cloud — the shared budget is still alive.
func TestFailoverLocalTimeoutFallsBack(t *testing.T) {
	local := &funcLLM{fn: func(ctx context.Context, _ LLMRequest) (*LLMResponse, error) {
		<-ctx.Done() // blocks until the leg deadline fires
		return nil, ctx.Err()
	}}
	cloud := &fakeLLM{text: "cloud answer"}
	c := newTestFailover(local, cloud, true)
	c.legTimeout = 10 * time.Millisecond

	resp, err := c.Complete(context.Background(), LLMRequest{})
	if err != nil || resp.Text != "cloud answer" {
		t.Fatalf("resp=(%v,%v), want the cloud fallback after the leg timeout", resp, err)
	}
	if c.isHealthy() {
		t.Fatal("a local leg timeout must trip the breaker")
	}
}

// TestFailoverEmptyLocalReplyRetriesOnCloud: a 2xx with empty content from the free
// local leg (thinking ate the token budget) is retried on the cloud instead of
// shipping the empty-completion ⚠️ path; the breaker stays closed (the backend is up).
func TestFailoverEmptyLocalReplyRetriesOnCloud(t *testing.T) {
	local := &fakeLLM{text: "   ", model: "qwen-local"}
	cloud := &fakeLLM{text: "cloud answer", model: "grok-x"}
	c := newTestFailover(local, cloud, true)

	resp, err := c.Complete(context.Background(), LLMRequest{})
	if err != nil || resp.Text != "cloud answer" {
		t.Fatalf("resp=(%v,%v), want the cloud retry on an empty local reply", resp, err)
	}
	if !c.isHealthy() {
		t.Fatal("an empty local reply is not an availability failure — no trip")
	}
}

// TestFailoverBothLegsDown: local and cloud both failing surfaces the cloud error
// (the caller refunds + reacts) — never a silent nil.
func TestFailoverBothLegsDown(t *testing.T) {
	local := &fakeLLM{err: errors.New("tunnel down")}
	cloudErr := errors.New("xai down")
	cloud := &fakeLLM{err: cloudErr}
	c := newTestFailover(local, cloud, true)

	_, err := c.Complete(context.Background(), LLMRequest{})
	if !errors.Is(err, cloudErr) {
		t.Fatalf("err = %v, want the cloud error surfaced", err)
	}
}

// TestFailoverBreakerRecovery is the anti-flapping state machine: a request-path trip
// needs tripProbesToClose consecutive probe successes to close; a plain probe-down
// re-closes on the next success; a mid-streak probe failure resets the streak.
func TestFailoverBreakerRecovery(t *testing.T) {
	c := newTestFailover(&fakeLLM{}, &fakeLLM{}, false)

	// Startup: one successful probe brings the local leg online.
	c.noteProbe(true)
	if !c.isHealthy() {
		t.Fatal("one probe success must close the breaker at startup")
	}

	// Request-path trip: one probe success is NOT enough, two are.
	c.trip()
	c.noteProbe(true)
	if c.isHealthy() {
		t.Fatal("after a trip one probe success must not re-close the breaker")
	}
	c.noteProbe(true)
	if !c.isHealthy() {
		t.Fatalf("after a trip %d probe successes must re-close the breaker", tripProbesToClose)
	}

	// A mid-streak failure resets the streak.
	c.trip()
	c.noteProbe(true)
	c.noteProbe(false)
	c.noteProbe(true)
	if c.isHealthy() {
		t.Fatal("a probe failure mid-streak must reset the recovery counter")
	}
	c.noteProbe(true)
	if !c.isHealthy() {
		t.Fatal("a full consecutive streak after the reset must re-close the breaker")
	}

	// Plain probe-down (no request trip): the next single success re-closes.
	c.noteProbe(false)
	if c.isHealthy() {
		t.Fatal("a probe failure must open the breaker")
	}
	c.noteProbe(true)
	if !c.isHealthy() {
		t.Fatal("plain probe-down must recover on one success")
	}
}

// TestFailoverNeverOnlineWarnsOnce: a local leg that has never come online since
// boot (a wrong LOCAL_LLM_BASE_URL 404s every probe) must WARN exactly once — an
// enabled-but-inoperative feature silently billing the cloud is the failure the
// fail-fast config comment promises to prevent — while a leg that WAS online logs
// its down transition at INFO only.
func TestFailoverNeverOnlineWarnsOnce(t *testing.T) {
	var buf bytes.Buffer
	c := newTestFailover(&fakeLLM{}, &fakeLLM{}, false)
	c.log = slog.New(slog.NewTextHandler(&buf, nil))

	for i := 0; i < neverOnlineWarnAfter-1; i++ {
		c.noteProbe(false)
	}
	if strings.Contains(buf.String(), "never come online") {
		t.Fatalf("WARN fired before %d consecutive failures: %s", neverOnlineWarnAfter, buf.String())
	}
	c.noteProbe(false)
	if n := strings.Count(buf.String(), "never come online"); n != 1 {
		t.Fatalf("never-online WARN count = %d, want exactly 1", n)
	}
	c.noteProbe(false) // still down — must not repeat
	if n := strings.Count(buf.String(), "never come online"); n != 1 {
		t.Fatalf("never-online WARN repeated: count = %d", n)
	}

	// Once the leg HAS been online, a later long outage is a transition (INFO), not
	// the never-online WARN.
	buf.Reset()
	c.noteProbe(true)
	for i := 0; i < neverOnlineWarnAfter+1; i++ {
		c.noteProbe(false)
	}
	if strings.Contains(buf.String(), "never come online") {
		t.Fatalf("never-online WARN fired for a leg that was online: %s", buf.String())
	}
}

// TestFailoverProbeLoop exercises the real constructor end to end: the background
// prober sees a live /models endpoint and brings the local leg online; the endpoint
// dying takes it offline again.
func TestFailoverProbeLoop(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path != "/models" {
			t.Errorf("probe hit %q, want /models", r.URL.Path)
		}
		w.WriteHeader(http.StatusOK)
	}))
	defer srv.Close()

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	cfg := &Config{
		LocalLLMBaseURL:        srv.URL,
		LocalLLMTimeout:        time.Second,
		LocalLLMHealthInterval: 10 * time.Millisecond,
	}
	c := NewFailoverLLMClient(ctx, &fakeLLM{text: "local"}, &fakeLLM{text: "cloud"}, cfg, discardLog()).(*failoverLLMClient)

	waitFor := func(want bool, what string) {
		t.Helper()
		deadline := time.Now().Add(3 * time.Second)
		for c.isHealthy() != want {
			if time.Now().After(deadline) {
				t.Fatalf("timed out waiting for %s", what)
			}
			time.Sleep(5 * time.Millisecond)
		}
	}
	waitFor(true, "the prober to bring the local leg online")
	srv.Close()
	waitFor(false, "the prober to take the dead endpoint offline")
}