easyai-ai-gateway/apps/api/internal/httpapi/availability_timeout_integration_test.go

package httpapi

import (
	"bytes"
	"context"
	"encoding/json"
	"io"
	"log/slog"
	"net/http"
	"net/http/httptest"
	"net/url"
	"os"
	"strings"
	"testing"
	"time"

	"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)

func TestReadyReturnsPostgresUnavailableWithinTwoSeconds(t *testing.T) {
	db := newExhaustedPostgresStore(t)
	server := &Server{store: db, logger: slog.New(slog.NewJSONHandler(io.Discard, nil))}
	requestContext, cancel := context.WithTimeout(context.Background(), 4*time.Second)
	defer cancel()
	request := httptest.NewRequest(http.MethodGet, "/readyz", nil).WithContext(requestContext)
	recorder := httptest.NewRecorder()

	startedAt := time.Now()
	server.ready(recorder, request)
	elapsed := time.Since(startedAt)

	assertUnavailableResponse(t, recorder, "POSTGRES_UNAVAILABLE", "postgres unavailable")
	if elapsed > 3*time.Second {
		t.Fatalf("readiness timeout took %s, want no more than 3s", elapsed)
	}
}

func TestReadyUsesReservedCriticalPoolWhenExecutionPoolIsBusy(t *testing.T) {
	executionDB := newExhaustedPostgresStore(t)
	criticalDB := newAvailablePostgresStore(t)
	server := &Server{
		store:             executionDB,
		coordinationStore: criticalDB,
		logger:            slog.New(slog.NewJSONHandler(io.Discard, nil)),
	}
	request := httptest.NewRequest(http.MethodGet, "/readyz", nil)
	recorder := httptest.NewRecorder()

	server.ready(recorder, request)

	if recorder.Code != http.StatusOK {
		t.Fatalf("readiness status=%d, want 200 while critical pool is available; body=%s", recorder.Code, recorder.Body.String())
	}
}

func TestLoginReturnsAuthStoreUnavailableWithinFiveSeconds(t *testing.T) {
	db := newExhaustedPostgresStore(t)
	var logs bytes.Buffer
	server := &Server{store: db, logger: slog.New(slog.NewJSONHandler(&logs, nil))}
	requestContext, cancel := context.WithTimeout(context.Background(), 7*time.Second)
	defer cancel()
	request := httptest.NewRequest(http.MethodPost, "/api/v1/auth/login", strings.NewReader(`{"account":"timeout-test-account","password":"timeout-test-password"}`)).WithContext(requestContext)
	recorder := httptest.NewRecorder()

	startedAt := time.Now()
	server.login(recorder, request)
	elapsed := time.Since(startedAt)

	assertUnavailableResponse(t, recorder, "AUTH_STORE_UNAVAILABLE", "authentication service temporarily unavailable")
	if elapsed > 6*time.Second {
		t.Fatalf("login timeout took %s, want no more than 6s", elapsed)
	}
	logOutput := logs.String()
	for _, field := range []string{"postgres_pool_max_connections", "postgres_pool_acquired_connections", "postgres_pool_idle_connections", "postgres_pool_empty_acquire_count", "postgres_pool_canceled_acquire_count"} {
		if !strings.Contains(logOutput, field) {
			t.Fatalf("login failure log did not include %q: %s", field, logOutput)
		}
	}
	if strings.Contains(logOutput, "timeout-test-account") || strings.Contains(logOutput, "timeout-test-password") {
		t.Fatalf("login failure log exposed credentials: %s", logOutput)
	}
}

func newExhaustedPostgresStore(t *testing.T) *store.Store {
	t.Helper()
	databaseURL := strings.TrimSpace(os.Getenv("AI_GATEWAY_TEST_DATABASE_URL"))
	if databaseURL == "" {
		t.Skip("set AI_GATEWAY_TEST_DATABASE_URL to run PostgreSQL availability timeout tests")
	}
	parsed, err := url.Parse(databaseURL)
	if err != nil {
		t.Fatalf("parse test database URL: %v", err)
	}
	query := parsed.Query()
	query.Set("pool_max_conns", "1")
	query.Set("pool_min_conns", "0")
	parsed.RawQuery = query.Encode()

	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	db, err := store.Connect(ctx, parsed.String())
	if err != nil {
		t.Fatalf("connect timeout test store: %v", err)
	}
	t.Cleanup(db.Close)
	connection, err := db.Pool().Acquire(ctx)
	if err != nil {
		t.Fatalf("exhaust timeout test pool: %v", err)
	}
	t.Cleanup(connection.Release)
	return db
}

func newAvailablePostgresStore(t *testing.T) *store.Store {
	t.Helper()
	databaseURL := strings.TrimSpace(os.Getenv("AI_GATEWAY_TEST_DATABASE_URL"))
	if databaseURL == "" {
		t.Skip("set AI_GATEWAY_TEST_DATABASE_URL to run PostgreSQL availability timeout tests")
	}
	parsed, err := url.Parse(databaseURL)
	if err != nil {
		t.Fatalf("parse test database URL: %v", err)
	}
	query := parsed.Query()
	query.Set("pool_max_conns", "1")
	query.Set("pool_min_conns", "0")
	parsed.RawQuery = query.Encode()

	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	db, err := store.Connect(ctx, parsed.String())
	if err != nil {
		t.Fatalf("connect available test store: %v", err)
	}
	t.Cleanup(db.Close)
	return db
}

func assertUnavailableResponse(t *testing.T, recorder *httptest.ResponseRecorder, expectedCode, expectedMessage string) {
	t.Helper()
	if recorder.Code != http.StatusServiceUnavailable {
		t.Fatalf("status = %d, want 503; body=%s", recorder.Code, recorder.Body.String())
	}
	var envelope ErrorEnvelope
	if err := json.Unmarshal(recorder.Body.Bytes(), &envelope); err != nil {
		t.Fatalf("decode unavailable response: %v", err)
	}
	if envelope.Error.Code != expectedCode {
		t.Fatalf("error code = %q, want %q; body=%s", envelope.Error.Code, expectedCode, recorder.Body.String())
	}
	if envelope.Error.Message != expectedMessage {
		t.Fatalf("error message = %q, want %q; body=%s", envelope.Error.Message, expectedMessage, recorder.Body.String())
	}
}