veylant/test/integration/e2e_proxy_test.go

//go:build integration

// Package integration contains Sprint 11 E2E tests (E11-01a).
// Batch 1 covers: auth flows, proxy chat, streaming SSE, audit logging,
// rate limiting, and feature flags (pii_enabled, billing_enabled).
//
// Run with: go test -tags integration -v ./test/integration/ -run TestE2E_
// All tests use httptest.Server + in-memory stubs — no external services needed.
package integration

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

	"github.com/go-chi/chi/v5"
	chimiddleware "github.com/go-chi/chi/v5/middleware"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"go.uber.org/zap"

	"github.com/veylant/ia-gateway/internal/apierror"
	"github.com/veylant/ia-gateway/internal/auditlog"
	"github.com/veylant/ia-gateway/internal/config"
	"github.com/veylant/ia-gateway/internal/flags"
	"github.com/veylant/ia-gateway/internal/health"
	"github.com/veylant/ia-gateway/internal/middleware"
	"github.com/veylant/ia-gateway/internal/provider"
	"github.com/veylant/ia-gateway/internal/proxy"
	"github.com/veylant/ia-gateway/internal/ratelimit"
	"github.com/veylant/ia-gateway/internal/router"
)

// ─── Test constants ──────────────────────────────────────────────────────────

const (
	e2eTenantID = "00000000-0000-0000-0000-000000000042"
	e2eUserID   = "user-e2e-test"
)

// ─── Helpers ──────────────────────────────────────────────────────────────────

// e2eAdminClaims returns authenticated admin claims for the E2E tenant.
func e2eAdminClaims() *middleware.UserClaims {
	return &middleware.UserClaims{
		UserID:   e2eUserID,
		TenantID: e2eTenantID,
		Email:    "e2e-admin@veylant.test",
		Roles:    []string{"admin"},
	}
}

// e2eUserClaims returns authenticated user-role claims for the E2E tenant.
func e2eUserClaims(roles ...string) *middleware.UserClaims {
	if len(roles) == 0 {
		roles = []string{"user"}
	}
	return &middleware.UserClaims{
		UserID:   e2eUserID,
		TenantID: e2eTenantID,
		Email:    "e2e-user@veylant.test",
		Roles:    roles,
	}
}

// proxyServerOptions holds optional components for building a proxy test server.
type proxyServerOptions struct {
	adapter   provider.Adapter
	flagStore flags.FlagStore
	auditLog  auditlog.Logger
	verifier  middleware.TokenVerifier
	// rateLimiter, if nil, a permissive default is used.
	rateLimiter *ratelimit.Limiter
}

// buildProxyServer creates a fully wired httptest.Server with chi router,
// auth middleware, rate limiter, and proxy handler. All components are
// in-memory; no external services are required.
func buildProxyServer(t *testing.T, opts proxyServerOptions) *httptest.Server {
	t.Helper()

	if opts.adapter == nil {
		opts.adapter = &e2eStubAdapter{
			resp: &provider.ChatResponse{
				ID:      "chatcmpl-e2e",
				Model:   "gpt-4o",
				Choices: []provider.Choice{{Index: 0, Message: provider.Message{Role: "assistant", Content: "Hello!"}}},
				Usage:   provider.Usage{PromptTokens: 10, CompletionTokens: 5, TotalTokens: 15},
			},
		}
	}
	if opts.flagStore == nil {
		opts.flagStore = flags.NewMemFlagStore()
	}
	if opts.verifier == nil {
		opts.verifier = &middleware.MockVerifier{Claims: e2eAdminClaims()}
	}
	if opts.rateLimiter == nil {
		opts.rateLimiter = ratelimit.New(ratelimit.RateLimitConfig{
			RequestsPerMin: 6000,
			BurstSize:      1000,
			UserRPM:        6000,
			UserBurst:      1000,
			IsEnabled:      true,
		}, zap.NewNop())
	}

	// Build a single-adapter router (no RBAC enforcement by default — admin role).
	rbac := &config.RBACConfig{
		UserAllowedModels:   []string{"gpt-4o", "gpt-4o-mini", "mistral-small"},
		AuditorCanComplete:  false,
	}
	providerRouter := router.New(
		map[string]provider.Adapter{"openai": opts.adapter},
		rbac,
		zap.NewNop(),
	)
	providerRouter.WithFlagStore(opts.flagStore)

	proxyHandler := proxy.NewWithAudit(providerRouter, zap.NewNop(), nil, opts.auditLog, nil).
		WithFlagStore(opts.flagStore)

	r := chi.NewRouter()
	r.Use(chimiddleware.Recoverer)
	r.Use(middleware.RequestID)
	r.Get("/healthz", health.Handler)

	r.Route("/v1", func(r chi.Router) {
		r.Use(middleware.Auth(opts.verifier))
		r.Use(middleware.RateLimit(opts.rateLimiter))
		r.Post("/chat/completions", proxyHandler.ServeHTTP)
	})

	srv := httptest.NewServer(r)
	t.Cleanup(srv.Close)
	return srv
}

// doJSON sends a JSON request to the test server and returns the response.
func doJSON(t *testing.T, client *http.Client, method, url, token string, body interface{}) *http.Response {
	t.Helper()
	var r io.Reader
	if body != nil {
		b, err := json.Marshal(body)
		require.NoError(t, err)
		r = bytes.NewReader(b)
	}
	req, err := http.NewRequestWithContext(context.Background(), method, url, r)
	require.NoError(t, err)
	if token != "" {
		req.Header.Set("Authorization", "Bearer "+token)
	}
	req.Header.Set("Content-Type", "application/json")
	resp, err := client.Do(req)
	require.NoError(t, err)
	return resp
}

// chatBody returns a minimal /v1/chat/completions request body.
func chatBody(model string, stream bool) map[string]interface{} {
	return map[string]interface{}{
		"model":    model,
		"messages": []map[string]string{{"role": "user", "content": "Dis bonjour."}},
		"stream":   stream,
	}
}

// ─── E2E stub adapter ─────────────────────────────────────────────────────────

// e2eStubAdapter is a minimal provider.Adapter for E2E tests.
type e2eStubAdapter struct {
	resp      *provider.ChatResponse
	sendErr   error
	sseLines  []string
}

func (s *e2eStubAdapter) Validate(req *provider.ChatRequest) error {
	if req.Model == "" {
		return apierror.NewBadRequestError("model required")
	}
	return nil
}

func (s *e2eStubAdapter) Send(_ context.Context, _ *provider.ChatRequest) (*provider.ChatResponse, error) {
	if s.sendErr != nil {
		return nil, s.sendErr
	}
	resp := *s.resp // copy
	return &resp, nil
}

func (s *e2eStubAdapter) Stream(_ context.Context, _ *provider.ChatRequest, w http.ResponseWriter) error {
	flusher, ok := w.(http.Flusher)
	if !ok {
		return fmt.Errorf("not flushable")
	}
	lines := s.sseLines
	if len(lines) == 0 {
		lines = []string{
			`data: {"id":"chatcmpl-e2e","object":"chat.completion.chunk","choices":[{"delta":{"content":"Hello"}}]}`,
			`data: [DONE]`,
		}
	}
	for _, line := range lines {
		fmt.Fprintf(w, "%s\n\n", line) //nolint:errcheck
		flusher.Flush()
	}
	return nil
}

func (s *e2eStubAdapter) HealthCheck(_ context.Context) error { return nil }

// ─── Scenario 1: GET /healthz ─────────────────────────────────────────────────

// TestE2E_HealthCheck_Returns200 verifies the health endpoint is reachable and
// returns status "ok". This is the first check in any deployment pipeline.
func TestE2E_HealthCheck_Returns200(t *testing.T) {
	srv := buildProxyServer(t, proxyServerOptions{})
	resp := doJSON(t, srv.Client(), http.MethodGet, srv.URL+"/healthz", "", nil)
	defer resp.Body.Close() //nolint:errcheck

	assert.Equal(t, http.StatusOK, resp.StatusCode)

	var body struct {
		Status string `json:"status"`
	}
	require.NoError(t, json.NewDecoder(resp.Body).Decode(&body))
	assert.Equal(t, "ok", body.Status)
}

// ─── Scenario 2: Auth — no token → 401 ───────────────────────────────────────

// TestE2E_Auth_NoToken_Returns401 verifies that unauthenticated requests are
// rejected with 401 before they reach the proxy handler.
func TestE2E_Auth_NoToken_Returns401(t *testing.T) {
	srv := buildProxyServer(t, proxyServerOptions{})
	resp := doJSON(t, srv.Client(), http.MethodPost, srv.URL+"/v1/chat/completions", "", chatBody("gpt-4o", false))
	defer resp.Body.Close() //nolint:errcheck

	assert.Equal(t, http.StatusUnauthorized, resp.StatusCode)

	var body struct {
		Error struct{ Type string } `json:"error"`
	}
	require.NoError(t, json.NewDecoder(resp.Body).Decode(&body))
	assert.Equal(t, "authentication_error", body.Error.Type)
}

// ─── Scenario 3: Auth — invalid token → 401 ──────────────────────────────────

// TestE2E_Auth_InvalidToken_Returns401 verifies that a structurally invalid token
// causes the OIDC verifier to return an error and the middleware to reply 401.
func TestE2E_Auth_InvalidToken_Returns401(t *testing.T) {
	srv := buildProxyServer(t, proxyServerOptions{
		verifier: &middleware.MockVerifier{Err: fmt.Errorf("token expired")},
	})
	resp := doJSON(t, srv.Client(), http.MethodPost, srv.URL+"/v1/chat/completions", "invalid-token", chatBody("gpt-4o", false))
	defer resp.Body.Close() //nolint:errcheck

	assert.Equal(t, http.StatusUnauthorized, resp.StatusCode)
}

// ─── Scenario 4: Auth — valid JWT → proxy responds 200 ───────────────────────

// TestE2E_Auth_ValidToken_ProxyResponds200 verifies end-to-end: valid JWT is
// accepted, proxy dispatches to the stub adapter, and returns a valid response.
func TestE2E_Auth_ValidToken_ProxyResponds200(t *testing.T) {
	srv := buildProxyServer(t, proxyServerOptions{})
	resp := doJSON(t, srv.Client(), http.MethodPost, srv.URL+"/v1/chat/completions", "valid-token", chatBody("gpt-4o", false))
	defer resp.Body.Close() //nolint:errcheck

	assert.Equal(t, http.StatusOK, resp.StatusCode)
	assert.Equal(t, "application/json", resp.Header.Get("Content-Type"))
}

// ─── Scenario 5: Basic chat — response has choices ───────────────────────────

// TestE2E_Proxy_BasicChat_HasChoices verifies that the LLM response is properly
// forwarded: the `choices` array is present and contains the stub response.
func TestE2E_Proxy_BasicChat_HasChoices(t *testing.T) {
	srv := buildProxyServer(t, proxyServerOptions{})
	resp := doJSON(t, srv.Client(), http.MethodPost, srv.URL+"/v1/chat/completions", "valid-token", chatBody("gpt-4o", false))
	defer resp.Body.Close() //nolint:errcheck

	require.Equal(t, http.StatusOK, resp.StatusCode)

	var body struct {
		Choices []struct {
			Message struct{ Content string } `json:"message"`
		} `json:"choices"`
	}
	require.NoError(t, json.NewDecoder(resp.Body).Decode(&body))
	require.Len(t, body.Choices, 1)
	assert.Equal(t, "Hello!", body.Choices[0].Message.Content)
}

// ─── Scenario 6: Streaming SSE ───────────────────────────────────────────────

// TestE2E_Proxy_Streaming_SSEFormat verifies that a stream:true request receives
// Content-Type: text/event-stream and at least one SSE data chunk.
func TestE2E_Proxy_Streaming_SSEFormat(t *testing.T) {
	srv := buildProxyServer(t, proxyServerOptions{})

	body, err := json.Marshal(chatBody("gpt-4o", true))
	require.NoError(t, err)

	req, err := http.NewRequestWithContext(context.Background(), http.MethodPost,
		srv.URL+"/v1/chat/completions", bytes.NewReader(body))
	require.NoError(t, err)
	req.Header.Set("Authorization", "Bearer valid-token")
	req.Header.Set("Content-Type", "application/json")

	resp, err := srv.Client().Do(req)
	require.NoError(t, err)
	defer resp.Body.Close() //nolint:errcheck

	assert.Equal(t, http.StatusOK, resp.StatusCode)
	assert.Equal(t, "text/event-stream", resp.Header.Get("Content-Type"))

	buf, err := io.ReadAll(resp.Body)
	require.NoError(t, err)
	assert.Contains(t, string(buf), "data:")
}

// ─── Scenario 7: Audit log entry created ─────────────────────────────────────

// TestE2E_Proxy_AuditLog_EntryCreated verifies that after a successful request,
// a single audit entry is recorded with the correct tenant and user fields.
func TestE2E_Proxy_AuditLog_EntryCreated(t *testing.T) {
	memLog := auditlog.NewMemLogger()

	srv := buildProxyServer(t, proxyServerOptions{auditLog: memLog})
	resp := doJSON(t, srv.Client(), http.MethodPost, srv.URL+"/v1/chat/completions", "valid-token", chatBody("gpt-4o", false))
	defer resp.Body.Close() //nolint:errcheck
	require.Equal(t, http.StatusOK, resp.StatusCode)

	// The audit logger is async but MemLogger.Log is synchronous.
	entries := memLog.Entries()
	require.Len(t, entries, 1)

	entry := entries[0]
	assert.Equal(t, e2eTenantID, entry.TenantID)
	assert.Equal(t, e2eUserID, entry.UserID)
	assert.Equal(t, "gpt-4o", entry.ModelRequested)
	assert.Equal(t, "ok", entry.Status)
	assert.NotEmpty(t, entry.PromptHash, "prompt hash should be recorded")
}

// ─── Scenario 8: Rate limiting → 429 ─────────────────────────────────────────

// TestE2E_Proxy_RateLimit_Returns429 verifies that a tenant with a burst of 1
// receives 429 on the second request within the same burst window.
func TestE2E_Proxy_RateLimit_Returns429(t *testing.T) {
	tightLimiter := ratelimit.New(ratelimit.RateLimitConfig{
		RequestsPerMin: 60,
		BurstSize:      1, // burst of 1 → second request is rate-limited
		UserRPM:        600,
		UserBurst:      100,
		IsEnabled:      true,
	}, zap.NewNop())

	srv := buildProxyServer(t, proxyServerOptions{rateLimiter: tightLimiter})

	client := srv.Client()

	// First request: should succeed.
	resp1 := doJSON(t, client, http.MethodPost, srv.URL+"/v1/chat/completions", "valid-token", chatBody("gpt-4o", false))
	defer resp1.Body.Close() //nolint:errcheck
	assert.Equal(t, http.StatusOK, resp1.StatusCode, "first request should be allowed")

	// Second request: burst exhausted → 429.
	resp2 := doJSON(t, client, http.MethodPost, srv.URL+"/v1/chat/completions", "valid-token", chatBody("gpt-4o", false))
	defer resp2.Body.Close() //nolint:errcheck
	assert.Equal(t, http.StatusTooManyRequests, resp2.StatusCode, "second request should be rate limited")

	var body struct {
		Error struct{ Type string } `json:"error"`
	}
	_ = json.NewDecoder(resp2.Body).Decode(&body)
	assert.Equal(t, "rate_limit_error", body.Error.Type)
}

// ─── Scenario 9: billing_enabled flag → CostUSD zeroed ───────────────────────

// TestE2E_Proxy_Billing_Flag_Disabled_ZerosCost verifies that when billing_enabled
// is set to false for a tenant, the CostUSD field in the audit entry is zeroed.
func TestE2E_Proxy_Billing_Flag_Disabled_ZerosCost(t *testing.T) {
	memLog := auditlog.NewMemLogger()
	fs := flags.NewMemFlagStore()

	// Explicitly disable billing for the test tenant.
	ctx := context.Background()
	_, err := fs.Set(ctx, e2eTenantID, "billing_enabled", false)
	require.NoError(t, err)
	// Also seed pii_enabled=true so the flag store is "warm" (like after migration 000009).
	_, err = fs.Set(ctx, "", "pii_enabled", true)
	require.NoError(t, err)

	srv := buildProxyServer(t, proxyServerOptions{auditLog: memLog, flagStore: fs})
	resp := doJSON(t, srv.Client(), http.MethodPost, srv.URL+"/v1/chat/completions", "valid-token", chatBody("gpt-4o", false))
	defer resp.Body.Close() //nolint:errcheck
	require.Equal(t, http.StatusOK, resp.StatusCode)

	entries := memLog.Entries()
	require.Len(t, entries, 1)
	assert.Equal(t, 0.0, entries[0].CostUSD, "billing_enabled=false should zero the cost")
}

// ─── Scenario 10: RBAC — user role + allowed model → 200 ─────────────────────

// TestE2E_Auth_RBAC_UserRole_AllowedModel_Returns200 verifies that a `user`-role
// JWT can call /v1/chat/completions for models in the allowlist (gpt-4o-mini).
func TestE2E_Auth_RBAC_UserRole_AllowedModel_Returns200(t *testing.T) {
	// User-role verifier (not admin).
	userVerifier := &middleware.MockVerifier{Claims: e2eUserClaims("user")}

	// Adapter that responds to gpt-4o-mini requests.
	adapter := &e2eStubAdapter{
		resp: &provider.ChatResponse{
			ID:      "chatcmpl-mini",
			Model:   "gpt-4o-mini",
			Choices: []provider.Choice{{Message: provider.Message{Role: "assistant", Content: "Hi!"}}},
			Usage:   provider.Usage{TotalTokens: 5},
		},
	}

	srv := buildProxyServer(t, proxyServerOptions{
		adapter:  adapter,
		verifier: userVerifier,
	})

	// gpt-4o-mini is in the default UserAllowedModels — should be allowed.
	resp := doJSON(t, srv.Client(), http.MethodPost, srv.URL+"/v1/chat/completions", "user-token", chatBody("gpt-4o-mini", false))
	defer resp.Body.Close() //nolint:errcheck

	assert.Equal(t, http.StatusOK, resp.StatusCode,
		"user role should be allowed to use gpt-4o-mini (in allowlist)")
}

// ─── Additional: request ID propagation ──────────────────────────────────────

// TestE2E_Proxy_RequestID_InLogs verifies that the X-Request-Id header is
// set on the request context and the request ID is stored in the audit entry.
func TestE2E_Proxy_RequestID_InLogs(t *testing.T) {
	memLog := auditlog.NewMemLogger()
	srv := buildProxyServer(t, proxyServerOptions{auditLog: memLog})

	req, err := http.NewRequestWithContext(context.Background(), http.MethodPost,
		srv.URL+"/v1/chat/completions",
		strings.NewReader(`{"model":"gpt-4o","messages":[{"role":"user","content":"hi"}]}`))
	require.NoError(t, err)
	req.Header.Set("Authorization", "Bearer valid-token")
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("X-Request-Id", "e2e-test-req-id-42")

	resp, err := srv.Client().Do(req)
	require.NoError(t, err)
	defer resp.Body.Close() //nolint:errcheck
	require.Equal(t, http.StatusOK, resp.StatusCode)

	entries := memLog.Entries()
	require.Len(t, entries, 1)
	assert.NotEmpty(t, entries[0].RequestID, "request ID should be recorded in audit entry")
}

// ─── Additional: pii_enabled flag does not break normal flow ─────────────────

// TestE2E_Proxy_PII_Flag_Enabled_NoClientIsNoOp verifies that pii_enabled=true
// with no PII client configured (nil) does not cause an error — the proxy
// degrades gracefully by skipping PII anonymization.
func TestE2E_Proxy_PII_Flag_Enabled_NoClientIsNoOp(t *testing.T) {
	fs := flags.NewMemFlagStore()
	ctx := context.Background()
	// Seed pii_enabled=true globally (as migration 000009 does).
	_, err := fs.Set(ctx, "", "pii_enabled", true)
	require.NoError(t, err)

	memLog := auditlog.NewMemLogger()
	srv := buildProxyServer(t, proxyServerOptions{flagStore: fs, auditLog: memLog})

	resp := doJSON(t, srv.Client(), http.MethodPost, srv.URL+"/v1/chat/completions", "valid-token",
		map[string]interface{}{
			"model":    "gpt-4o",
			"messages": []map[string]string{{"role": "user", "content": "Mon IBAN est FR76 3000 6000 0112 3456 7890 189"}},
		})
	defer resp.Body.Close() //nolint:errcheck

	// pii_enabled=true but no PII client → still succeeds, no error.
	assert.Equal(t, http.StatusOK, resp.StatusCode)
	entries := memLog.Entries()
	require.Len(t, entries, 1)
	// No PII client → entity count is 0.
	assert.Equal(t, 0, entries[0].PIIEntityCount)
}

// ─── Additional: invalid JSON body → 400 ─────────────────────────────────────

// TestE2E_Proxy_InvalidJSON_Returns400 verifies that a malformed JSON body
// is rejected immediately with 400 invalid_request_error.
func TestE2E_Proxy_InvalidJSON_Returns400(t *testing.T) {
	srv := buildProxyServer(t, proxyServerOptions{})

	req, err := http.NewRequestWithContext(context.Background(), http.MethodPost,
		srv.URL+"/v1/chat/completions",
		strings.NewReader(`{not valid json`))
	require.NoError(t, err)
	req.Header.Set("Authorization", "Bearer valid-token")
	req.Header.Set("Content-Type", "application/json")

	resp, err := srv.Client().Do(req)
	require.NoError(t, err)
	defer resp.Body.Close() //nolint:errcheck

	assert.Equal(t, http.StatusBadRequest, resp.StatusCode)

	var body struct {
		Error struct{ Type string } `json:"error"`
	}
	require.NoError(t, json.NewDecoder(resp.Body).Decode(&body))
	assert.Equal(t, "invalid_request_error", body.Error.Type)
}

// ─── Additional: upstream error → 502 ────────────────────────────────────────

// TestE2E_Proxy_UpstreamError_Returns502 verifies that when the LLM provider
// returns an error, the proxy responds with 502 and the correct error type.
func TestE2E_Proxy_UpstreamError_Returns502(t *testing.T) {
	errAdapter := &e2eStubAdapter{
		resp:    nil,
		sendErr: fmt.Errorf("connection refused"),
	}
	srv := buildProxyServer(t, proxyServerOptions{adapter: errAdapter})

	resp := doJSON(t, srv.Client(), http.MethodPost, srv.URL+"/v1/chat/completions", "valid-token", chatBody("gpt-4o", false))
	defer resp.Body.Close() //nolint:errcheck

	assert.Equal(t, http.StatusBadGateway, resp.StatusCode)
}

// ─── Additional: routing_enabled flag (via MemFlagStore seeding) ─────────────

// TestE2E_Proxy_RoutingFlag_Enabled_DefaultsBehavior verifies that explicitly
// setting routing_enabled=true (global default from migration 000009) does not
// break the static routing path when no routing engine is configured.
func TestE2E_Proxy_RoutingFlag_Enabled_DefaultsBehavior(t *testing.T) {
	fs := flags.NewMemFlagStore()
	ctx := context.Background()
	_, err := fs.Set(ctx, "", "routing_enabled", true)
	require.NoError(t, err)

	srv := buildProxyServer(t, proxyServerOptions{flagStore: fs})
	resp := doJSON(t, srv.Client(), http.MethodPost, srv.URL+"/v1/chat/completions", "valid-token", chatBody("gpt-4o", false))
	defer resp.Body.Close() //nolint:errcheck

	// routing_enabled=true + no engine → static prefix rules → should still work.
	assert.Equal(t, http.StatusOK, resp.StatusCode)
}

// ─── Ensure test timeout is reasonable ───────────────────────────────────────

// TestE2E_Proxy_BatchRunsUnder30s is a meta-test verifying that all proxy E2E
// tests complete well within the 10-minute CI budget. It times the package.
func TestE2E_Proxy_BatchRunsUnder30s(t *testing.T) {
	start := time.Now()
	t.Cleanup(func() {
		elapsed := time.Since(start)
		t.Logf("E2E proxy batch wall time: %v", elapsed)
		assert.Less(t, elapsed, 30*time.Second, "proxy E2E tests should complete in < 30s")
	})
}