AIRouter/test/integration/e2e_test.go

package integration

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

	"github.com/edgeai/gateway/internal/auth"
	"github.com/edgeai/gateway/internal/config"
	"github.com/edgeai/gateway/internal/observability"
	"github.com/edgeai/gateway/internal/server"
	"github.com/edgeai/gateway/pkg/api"
)

var testServer *httptest.Server

func TestMain(m *testing.M) {
	// Create a temp config
	cfg := &config.Config{
		Server: config.ServerConfig{
			Host: "127.0.0.1", Port: 0, AdminPort: 0, MaxRequestBodyMB: 5,
		},
		Auth: config.AuthConfig{Enabled: true, Methods: []string{"api_key"}},
		Scheduler: config.SchedulerConfig{
			MaxRunningTasks: 2, MaxQueuedTasks: 10, Fairness: "weighted_fair_queue",
			PriorityAgingSeconds: 5, ReservedRealtimeSlots: 1,
		},
		Timeouts: config.TimeoutConfig{
			DefaultConnectMs: 2000, DefaultQueueMs: 2000, DefaultFirstTokenMs: 5000,
			DefaultInferenceMs: 10000, DefaultIdleMs: 5000, DefaultTotalMs: 15000,
			CancelGracePeriodMs: 1000,
		},
		Context: config.ContextConfig{
			SafetyMarginRatio: 0.1, DefaultPolicy: "recent_only",
			MaxSessionMessages: 20, SessionIdleTTLMinutes: 5,
		},
		Models: map[string]config.ModelConfig{
			"test-chat": {
				Provider: "ollama", ActualModel: "qwen2.5:0.5b",
				Endpoint: "http://127.0.0.1:11434", ContextWindow: 4096,
				MaxOutputTokens: 256, MaxConcurrency: 1, Residency: "always",
				CancelSupported: true,
			},
		},
		Observability: config.ObservabilityConfig{
			MetricsEnabled: true, MetricsPath: "/metrics",
			PromptLogging: "metadata_only", LogLevel: "debug",
		},
		Storage: config.StorageConfig{
			SessionDB: "sqlite:///tmp/edgeai-int-test/sessions.db",
			TaskState: "sqlite:///tmp/edgeai-int-test/tasks.db",
		},
	}

	os.MkdirAll("/tmp/edgeai-int-test", 0755)
	defer os.RemoveAll("/tmp/edgeai-int-test")

	logger := observability.NewLogger(observability.LevelDebug, os.Stdout, "metadata_only")
	srv, err := server.New(cfg, logger)
	if err != nil {
		panic("failed to create test server: " + err.Error())
	}

	// Add a test API key for authenticated tests
	srv.Authenticator().AddKey("test-key", &auth.AppIdentity{
		AppID:         "test-app",
		TenantID:      "test-tenant",
		Name:          "test",
		AllowedModels: []string{}, // empty = all models
		IsAdmin:       true,
	})

	testServer = httptest.NewServer(srv.HTTPSrv.Handler)
	defer testServer.Close()

	m.Run()
}

func doRequest(t *testing.T, method, path string, body any, apiKey string) (*http.Response, []byte) {
	t.Helper()
	var buf bytes.Buffer
	if body != nil {
		json.NewEncoder(&buf).Encode(body)
	}
	req, _ := http.NewRequest(method, testServer.URL+path, &buf)
	req.Header.Set("Content-Type", "application/json")
	if apiKey != "" {
		req.Header.Set("Authorization", "Bearer "+apiKey)
	}
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		t.Fatalf("request failed: %v", err)
	}
	respBody := make([]byte, 0)
	if resp.Body != nil {
		buf := bytes.Buffer{}
		buf.ReadFrom(resp.Body)
		respBody = buf.Bytes()
		resp.Body.Close()
	}
	return resp, respBody
}

// E2E-001: Health check returns 200
func TestE2E001_HealthCheck(t *testing.T) {
	resp, body := doRequest(t, "GET", "/health", nil, "")
	if resp.StatusCode != 200 {
		t.Errorf("expected 200, got %d", resp.StatusCode)
	}
	var result map[string]string
	json.Unmarshal(body, &result)
	if result["status"] != "ok" {
		t.Errorf("expected status ok, got %s", result["status"])
	}
}

// E2E-002: Ready check returns 200 or 503
func TestE2E002_ReadyCheck(t *testing.T) {
	resp, _ := doRequest(t, "GET", "/ready", nil, "")
	if resp.StatusCode != 200 && resp.StatusCode != 503 {
		t.Errorf("expected 200 or 503, got %d", resp.StatusCode)
	}
}

// E2E-003: Metrics endpoint returns 200
func TestE2E003_MetricsEndpoint(t *testing.T) {
	resp, body := doRequest(t, "GET", "/metrics", nil, "")
	if resp.StatusCode != 200 {
		t.Errorf("expected 200, got %d", resp.StatusCode)
	}
	if !strings.Contains(string(body), "edgeai_") {
		t.Error("expected edgeai_ metrics in response")
	}
}

// E2E-004: Unauthenticated request returns 401
func TestE2E004_UnauthenticatedRequest(t *testing.T) {
	resp, _ := doRequest(t, "GET", "/v1/models", nil, "")
	if resp.StatusCode != 401 {
		t.Errorf("expected 401, got %d", resp.StatusCode)
	}
}

// E2E-005: Invalid API key returns 401
func TestE2E005_InvalidAPIKey(t *testing.T) {
	resp, body := doRequest(t, "GET", "/v1/models", nil, "invalid-key")
	if resp.StatusCode != 401 {
		t.Errorf("expected 401, got %d", resp.StatusCode)
	}
	var errResp map[string]any
	json.Unmarshal(body, &errResp)
	errBody := errResp["error"].(map[string]any)
	if errBody["code"] != "AUTH_FAILED" {
		t.Errorf("expected AUTH_FAILED, got %v", errBody["code"])
	}
}

// E2E-006: Missing Authorization header returns 401
func TestE2E006_MissingAuthHeader(t *testing.T) {
	req, _ := http.NewRequest("GET", testServer.URL+"/v1/models", nil)
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	if resp.StatusCode != 401 {
		t.Errorf("expected 401, got %d", resp.StatusCode)
	}
	resp.Body.Close()
}

// E2E-007: Malformed Authorization header returns 401
func TestE2E007_MalformedAuth(t *testing.T) {
	req, _ := http.NewRequest("GET", testServer.URL+"/v1/models", nil)
	req.Header.Set("Authorization", "Basic abc123")
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	if resp.StatusCode != 401 {
		t.Errorf("expected 401, got %d", resp.StatusCode)
	}
	resp.Body.Close()
}

// E2E-008: Chat completions with missing model returns 400
func TestE2E008_ChatMissingModel(t *testing.T) {
	resp, body := doRequest(t, "POST", "/v1/chat/completions", api.ChatRequest{
		Messages: []api.Message{{Role: "user", Content: "hello"}},
	}, "test-key")
	if resp.StatusCode != 400 {
		t.Errorf("expected 400, got %d", resp.StatusCode)
	}
	var errResp map[string]any
	json.Unmarshal(body, &errResp)
	errBody := errResp["error"].(map[string]any)
	if errBody["code"] != "INVALID_REQUEST" {
		t.Errorf("expected INVALID_REQUEST, got %v", errBody["code"])
	}
}

// E2E-009: Chat completions with missing messages returns 400
func TestE2E009_ChatMissingMessages(t *testing.T) {
	resp, _ := doRequest(t, "POST", "/v1/chat/completions", api.ChatRequest{
		Model: "test-chat",
	}, "test-key")
	if resp.StatusCode != 400 {
		t.Errorf("expected 400, got %d", resp.StatusCode)
	}
}

// E2E-010: Chat completions with unknown model returns 503
func TestE2E010_ChatUnknownModel(t *testing.T) {
	resp, body := doRequest(t, "POST", "/v1/chat/completions", api.ChatRequest{
		Model:    "nonexistent-model",
		Messages: []api.Message{{Role: "user", Content: "hello"}},
	}, "test-key")
	if resp.StatusCode != 503 {
		t.Errorf("expected 503, got %d", resp.StatusCode)
	}
	var errResp map[string]any
	json.Unmarshal(body, &errResp)
	errBody := errResp["error"].(map[string]any)
	if errBody["code"] != "MODEL_UNAVAILABLE" {
		t.Errorf("expected MODEL_UNAVAILABLE, got %v", errBody["code"])
	}
}

// E2E-011: Session creation returns 201
func TestE2E011_CreateSession(t *testing.T) {
	resp, body := doRequest(t, "POST", "/v1/sessions", api.SessionRequest{
		ApplicationID: "test-app",
		UserID:        "test-user",
	}, "test-key")
	if resp.StatusCode != 201 {
		t.Errorf("expected 201, got %d", resp.StatusCode)
	}
	var sessResp api.SessionResponse
	json.Unmarshal(body, &sessResp)
	if sessResp.SessionID == "" {
		t.Error("expected non-empty session ID")
	}
	if sessResp.ApplicationID != "test-app" {
		t.Errorf("expected app test-app, got %s", sessResp.ApplicationID)
	}
}

// E2E-012: Session creation without application_id returns 400
func TestE2E012_SessionMissingAppID(t *testing.T) {
	resp, _ := doRequest(t, "POST", "/v1/sessions", api.SessionRequest{}, "test-key")
	if resp.StatusCode != 400 {
		t.Errorf("expected 400, got %d", resp.StatusCode)
	}
}

// E2E-013: Request ID is set in response header
func TestE2E013_RequestIDHeader(t *testing.T) {
	req, _ := http.NewRequest("GET", testServer.URL+"/health", nil)
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	defer resp.Body.Close()

	requestID := resp.Header.Get("X-Request-ID")
	if requestID == "" {
		t.Error("expected X-Request-ID header to be set")
	}
}

// E2E-014: Custom request ID is preserved
func TestE2E014_CustomRequestID(t *testing.T) {
	customID := "my-custom-request-id-12345"
	req, _ := http.NewRequest("GET", testServer.URL+"/health", nil)
	req.Header.Set("X-Request-ID", customID)
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	defer resp.Body.Close()

	if resp.Header.Get("X-Request-ID") != customID {
		t.Errorf("expected %s, got %s", customID, resp.Header.Get("X-Request-ID"))
	}
}

// E2E-015: Error response contains request_id
func TestE2E015_ErrorContainsRequestID(t *testing.T) {
	resp, body := doRequest(t, "GET", "/v1/models", nil, "invalid-key")
	if resp.StatusCode != 401 {
		t.Fatalf("expected 401, got %d", resp.StatusCode)
	}
	var errResp map[string]any
	json.Unmarshal(body, &errResp)
	errBody := errResp["error"].(map[string]any)
	if errBody["request_id"] == nil || errBody["request_id"] == "" {
		t.Error("expected request_id in error response")
	}
}

// E2E-016: Body size limit is enforced
func TestE2E016_BodySizeLimit(t *testing.T) {
	largeContent := strings.Repeat("x", 6*1024*1024) // 6MB > 5MB limit
	req, _ := http.NewRequest("POST", testServer.URL+"/v1/chat/completions", strings.NewReader(largeContent))
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", "Bearer test-key")
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusBadRequest {
		t.Errorf("expected 400 for oversized body, got %d", resp.StatusCode)
	}
}

// E2E-017: Wrong HTTP method returns error
func TestE2E017_WrongMethod(t *testing.T) {
	resp, _ := doRequest(t, "DELETE", "/v1/chat/completions", nil, "test-key")
	if resp.StatusCode == 200 {
		t.Error("expected non-200 for DELETE on chat completions")
	}
}

// E2E-018: Concurrent requests don't crash the server
func TestE2E018_ConcurrentRequests(t *testing.T) {
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	done := make(chan error, 10)
	for i := 0; i < 10; i++ {
		go func(idx int) {
			resp, _ := doRequest(t, "GET", "/health", nil, "")
			if resp.StatusCode != 200 {
				done <- fmt.Errorf("goroutine %d: expected 200, got %d", idx, resp.StatusCode)
				return
			}
			done <- nil
		}(i)
	}

	for i := 0; i < 10; i++ {
		select {
		case err := <-done:
			if err != nil {
				t.Error(err)
			}
		case <-ctx.Done():
			t.Fatal("timeout waiting for concurrent requests")
		}
	}
}