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()) } }