goods/api/internal/handler/middleware_db_test.go

package handler

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

	"github.com/jackc/pgx/v5/pgxpool"

	"github.com/baicai2026-baicai/goods/api/internal/apikey"
	"github.com/baicai2026-baicai/goods/api/internal/ratelimit"
	"github.com/baicai2026-baicai/goods/api/internal/store"
)

// newRateLimitedHandler builds a handler backed by the test DB and a live Redis
// limiter, plus a freshly issued API key with the given per-minute limit. It
// skips when either backend is unavailable.
func newRateLimitedHandler(t *testing.T, keyLimit int) (h *Handler, plaintextKey string) {
	t.Helper()
	dsn := os.Getenv("OPENGOODS_DATABASE_URL")
	if dsn == "" {
		dsn = "postgres://opengoods:opengoods@localhost:5432/opengoods?sslmode=disable"
	}
	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
	defer cancel()
	pool, err := pgxpool.New(ctx, dsn)
	if err != nil {
		t.Skipf("no database: %v", err)
	}
	if err := pool.Ping(ctx); err != nil {
		pool.Close()
		t.Skipf("database not reachable: %v", err)
	}
	var hasTable bool
	if err := pool.QueryRow(ctx, "SELECT to_regclass('public.api_key') IS NOT NULL").Scan(&hasTable); err != nil || !hasTable {
		pool.Close()
		t.Skip("migrations not applied (api_key missing)")
	}

	redisURL := os.Getenv("OPENGOODS_REDIS_URL")
	if redisURL == "" {
		redisURL = "redis://localhost:6379/0"
	}
	limiter := ratelimit.New(redisURL)
	pingCtx, pingCancel := context.WithTimeout(context.Background(), time.Second)
	defer pingCancel()
	if err := limiter.Ping(pingCtx); err != nil {
		pool.Close()
		t.Skipf("redis not reachable: %v", err)
	}

	key, hash, prefix, err := apikey.Generate()
	if err != nil {
		pool.Close()
		t.Fatal(err)
	}
	name := fmt.Sprintf("test-key-%d", time.Now().UnixNano())
	if _, err := pool.Exec(context.Background(),
		`INSERT INTO api_key (name, key_prefix, key_hash, rate_limit_per_min) VALUES ($1,$2,$3,$4)`,
		name, prefix, hash, keyLimit); err != nil {
		pool.Close()
		t.Fatalf("insert api_key: %v", err)
	}
	t.Cleanup(func() {
		_, _ = pool.Exec(context.Background(), "DELETE FROM api_key WHERE key_hash=$1", hash)
		pool.Close()
	})

	return New(store.New(pool), nil).WithRateLimit(limiter, 60), key
}

func TestRateLimitHeadersAndKeyAuth(t *testing.T) {
	h, key := newRateLimitedHandler(t, 100)
	req := httptest.NewRequest(http.MethodGet, "/api/"+APIVersion+"/categories", nil)
	req.Header.Set("X-API-Key", key)
	rec := httptest.NewRecorder()
	h.Router().ServeHTTP(rec, req)

	if rec.Code != http.StatusOK {
		t.Fatalf("status = %d, body = %s", rec.Code, rec.Body.String())
	}
	if got := rec.Header().Get("X-RateLimit-Limit"); got != "100" {
		t.Fatalf("X-RateLimit-Limit = %q, want 100 (key limit)", got)
	}
	if rec.Header().Get("X-RateLimit-Remaining") == "" {
		t.Fatal("missing X-RateLimit-Remaining header")
	}
}

func TestInvalidKeyRejected(t *testing.T) {
	h, _ := newRateLimitedHandler(t, 100)
	req := httptest.NewRequest(http.MethodGet, "/api/"+APIVersion+"/categories", nil)
	req.Header.Set("X-API-Key", "og_live_thiskeydoesnotexist123456")
	rec := httptest.NewRecorder()
	h.Router().ServeHTTP(rec, req)
	if rec.Code != http.StatusUnauthorized {
		t.Fatalf("status = %d, want 401; body = %s", rec.Code, rec.Body.String())
	}
}

func TestRateLimitExceeded(t *testing.T) {
	h, key := newRateLimitedHandler(t, 1)
	do := func() int {
		req := httptest.NewRequest(http.MethodGet, "/api/"+APIVersion+"/categories", nil)
		req.Header.Set("X-API-Key", key)
		rec := httptest.NewRecorder()
		h.Router().ServeHTTP(rec, req)
		return rec.Code
	}
	if code := do(); code != http.StatusOK {
		t.Fatalf("first request status = %d, want 200", code)
	}
	if code := do(); code != http.StatusTooManyRequests {
		t.Fatalf("second request status = %d, want 429", code)
	}
}