goods/api/internal/handler/account_db_test.go

package handler

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

	"github.com/jackc/pgx/v5/pgxpool"
	"github.com/redis/go-redis/v9"

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

func cleanupCounter(t *testing.T, subject string) {
	t.Helper()
	redisURL := os.Getenv("OPENGOODS_REDIS_URL")
	if redisURL == "" {
		redisURL = "redis://localhost:6379/0"
	}
	opt, err := redis.ParseURL(redisURL)
	if err != nil {
		return
	}
	rdb := redis.NewClient(opt)
	defer rdb.Close()
	rdb.Del(context.Background(), "usage:total:"+subject)
}

// newQuotaHandler builds a handler backed by the test DB and a live Redis
// limiter, with a small anonymous quota so exhaustion is cheap to exercise.
func newQuotaHandler(t *testing.T, anonQuota int) *Handler {
	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 hasUser bool
	if err := pool.QueryRow(ctx, "SELECT to_regclass('public.app_user') IS NOT NULL").Scan(&hasUser); err != nil || !hasUser {
		pool.Close()
		t.Skip("migrations not applied")
	}
	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)
	}
	t.Cleanup(pool.Close)
	return New(store.New(pool), nil).
		WithRateLimit(limiter, 1000).
		WithQuotas(anonQuota, 300, 100000)
}

func cleanupAccount(t *testing.T, email string) {
	t.Helper()
	dsn := os.Getenv("OPENGOODS_DATABASE_URL")
	if dsn == "" {
		dsn = "postgres://opengoods:opengoods@localhost:5432/opengoods?sslmode=disable"
	}
	ctx := context.Background()
	pool, err := pgxpool.New(ctx, dsn)
	if err != nil {
		return
	}
	defer pool.Close()
	_, _ = pool.Exec(ctx, "DELETE FROM app_user WHERE lower(email)=lower($1)", email)
	_, _ = pool.Exec(ctx, "DELETE FROM api_key WHERE owner_email=$1", email)
}

func TestAnonTotalQuotaExhausts(t *testing.T) {
	h := newQuotaHandler(t, 3)
	// Unique client IP so the lifetime counter starts fresh for this test; the
	// counter never expires, so drop it afterwards to keep runs independent.
	n := time.Now().UnixNano()
	ip := fmt.Sprintf("203.%d.%d.%d", n/65536%256, n/256%256, n%256)
	t.Cleanup(func() { cleanupCounter(t, "ip:"+ip) })
	call := func() *httptest.ResponseRecorder {
		req := httptest.NewRequest(http.MethodGet, "/api/"+APIVersion+"/stats", nil)
		req.RemoteAddr = ip + ":12345"
		rec := httptest.NewRecorder()
		h.Router().ServeHTTP(rec, req)
		return rec
	}
	for i := 1; i <= 3; i++ {
		if rec := call(); rec.Code != http.StatusOK {
			t.Fatalf("call %d should be allowed, got %d (%s)", i, rec.Code, rec.Body.String())
		}
	}
	rec := call()
	if rec.Code != http.StatusForbidden {
		t.Fatalf("4th call should be 403 quota_exhausted, got %d (%s)", rec.Code, rec.Body.String())
	}
	if !strings.Contains(rec.Body.String(), "quota_exhausted") {
		t.Fatalf("expected quota_exhausted error, got %s", rec.Body.String())
	}
}

func TestRegisterIssuesHigherQuotaKey(t *testing.T) {
	h := newQuotaHandler(t, 1000)
	email := fmt.Sprintf("h-user-%d@example.com", time.Now().UnixNano())
	t.Cleanup(func() { cleanupAccount(t, email) })

	body := fmt.Sprintf(`{"email":%q,"password":"supersecret"}`, email)
	req := httptest.NewRequest(http.MethodPost, "/api/"+APIVersion+"/register", strings.NewReader(body))
	req.RemoteAddr = "198.51.100.7:9999"
	rec := httptest.NewRecorder()
	h.Router().ServeHTTP(rec, req)
	if rec.Code != http.StatusCreated {
		t.Fatalf("register status = %d (%s)", rec.Code, rec.Body.String())
	}
	var resp keyResponse
	if err := json.NewDecoder(rec.Body).Decode(&resp); err != nil {
		t.Fatal(err)
	}
	if resp.APIKey == "" || resp.QuotaTotal != 100000 || resp.RateLimitPerMin != 300 {
		t.Fatalf("unexpected register response: %+v", resp)
	}

	// A second registration with the same email conflicts.
	req2 := httptest.NewRequest(http.MethodPost, "/api/"+APIVersion+"/register", strings.NewReader(body))
	req2.RemoteAddr = "198.51.100.7:9999"
	rec2 := httptest.NewRecorder()
	h.Router().ServeHTTP(rec2, req2)
	if rec2.Code != http.StatusConflict {
		t.Fatalf("duplicate register status = %d (%s)", rec2.Code, rec2.Body.String())
	}
}