hyperguild/ingestion/internal/llm/router_test.go

package llm

import (
	"context"
	"net/http"
	"net/http/httptest"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestRouter_PrimarySucceeds(t *testing.T) {
	primary := mockServer(t, "from-primary")
	defer primary.Close()
	fallback := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		t.Error("fallback must not be called when primary succeeds")
	}))
	defer fallback.Close()

	r := &Router{
		Primary:  New(primary.URL, "", "m", time.Second),
		Fallback: New(fallback.URL, "", "m", time.Second),
	}
	out, err := r.Complete(context.Background(), "sys", "user")
	require.NoError(t, err)
	assert.Equal(t, "from-primary", out)
}

func TestRouter_FallsBackOnPrimaryError(t *testing.T) {
	primary := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		http.Error(w, "unavailable", http.StatusServiceUnavailable)
	}))
	defer primary.Close()
	fallback := mockServer(t, "from-fallback")
	defer fallback.Close()

	r := &Router{
		Primary:  New(primary.URL, "", "m", time.Second),
		Fallback: New(fallback.URL, "", "m", time.Second),
	}
	out, err := r.Complete(context.Background(), "sys", "user")
	require.NoError(t, err)
	assert.Equal(t, "from-fallback", out)
}

func TestRouter_BothFail(t *testing.T) {
	fail := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		http.Error(w, "err", http.StatusBadGateway)
	}))
	defer fail.Close()

	r := &Router{
		Primary:  New(fail.URL, "", "m", time.Second),
		Fallback: New(fail.URL, "", "m", time.Second),
	}
	_, err := r.Complete(context.Background(), "sys", "user")
	assert.Error(t, err)
}

func TestRouter_NilFallback(t *testing.T) {
	fail := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		http.Error(w, "err", http.StatusBadGateway)
	}))
	defer fail.Close()

	r := &Router{Primary: New(fail.URL, "", "m", time.Second)}
	_, err := r.Complete(context.Background(), "sys", "user")
	assert.Error(t, err)
}

func TestRouter_FallsBackToTertiaryWhenPrimaryAndFallbackFail(t *testing.T) {
	fail := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		http.Error(w, "unavailable", http.StatusServiceUnavailable)
	}))
	defer fail.Close()
	tertiary := mockServer(t, "from-tertiary")
	defer tertiary.Close()

	r := &Router{
		Primary:  New(fail.URL, "", "m", time.Second),
		Fallback: New(fail.URL, "", "m", time.Second),
		Tertiary: New(tertiary.URL, "", "m", time.Second),
	}
	out, err := r.Complete(context.Background(), "sys", "user")
	require.NoError(t, err)
	assert.Equal(t, "from-tertiary", out)
}

func TestRouter_TertiaryNotCalledWhenFallbackSucceeds(t *testing.T) {
	fail := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		http.Error(w, "unavailable", http.StatusServiceUnavailable)
	}))
	defer fail.Close()
	fallback := mockServer(t, "from-fallback")
	defer fallback.Close()
	tertiary := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		t.Error("tertiary must not be called when fallback succeeds")
	}))
	defer tertiary.Close()

	r := &Router{
		Primary:  New(fail.URL, "", "m", time.Second),
		Fallback: New(fallback.URL, "", "m", time.Second),
		Tertiary: New(tertiary.URL, "", "m", time.Second),
	}
	out, err := r.Complete(context.Background(), "sys", "user")
	require.NoError(t, err)
	assert.Equal(t, "from-fallback", out)
}

func TestRouter_AllThreeFail(t *testing.T) {
	fail := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		http.Error(w, "err", http.StatusBadGateway)
	}))
	defer fail.Close()

	r := &Router{
		Primary:  New(fail.URL, "", "m", time.Second),
		Fallback: New(fail.URL, "", "m", time.Second),
		Tertiary: New(fail.URL, "", "m", time.Second),
	}
	_, err := r.Complete(context.Background(), "sys", "user")
	assert.Error(t, err)
	assert.Contains(t, err.Error(), "primary llm")
	assert.Contains(t, err.Error(), "fallback llm")
	assert.Contains(t, err.Error(), "tertiary llm")
}

func TestRouter_NilTertiary(t *testing.T) {
	// Existing 2-tier deployments (no BRAIN_LLM_TERTIARY_URL) must keep working
	// exactly as before — Tertiary is opt-in, not required.
	fail := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		http.Error(w, "err", http.StatusBadGateway)
	}))
	defer fail.Close()

	r := &Router{
		Primary:  New(fail.URL, "", "m", time.Second),
		Fallback: New(fail.URL, "", "m", time.Second),
	}
	_, err := r.Complete(context.Background(), "sys", "user")
	assert.Error(t, err)
	assert.NotContains(t, err.Error(), "tertiary")
}