206 lines
7.0 KiB
Go
206 lines
7.0 KiB
Go
|
package ratelimit
|
|||
|
|
|
|||
|
|
import (
|
|||
|
|
"context"
|
|||
|
|
"os"
|
|||
|
|
"sync"
|
|||
|
|
"testing"
|
|||
|
|
"time"
|
|||
|
|
|
|||
|
|
"github.com/google/uuid"
|
|||
|
|
"github.com/jackc/pgx/v5"
|
|||
|
|
"github.com/jackc/pgx/v5/pgconn"
|
|||
|
|
"github.com/jackc/pgx/v5/pgxpool"
|
|||
|
|
"github.com/joho/godotenv"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
// newTestQuerier — паттерн из repository: транзакция поверх DATABASE_URL
|
|||
|
|
// с откатом в конце теста — проверяем на реальной БД, данные не остаются.
|
|||
|
|
// Без DATABASE_URL (или недоступной базы) тест пропускается.
|
|||
|
|
func newTestQuerier(t *testing.T) (Querier, context.Context) {
|
|||
|
|
t.Helper()
|
|||
|
|
|
|||
|
|
_ = godotenv.Load("../../.env")
|
|||
|
|
dsn := os.Getenv("DATABASE_URL")
|
|||
|
|
if dsn == "" {
|
|||
|
|
t.Skip("DATABASE_URL is not set — integration test skipped")
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
ctx := context.Background()
|
|||
|
|
pool, err := pgxpool.New(ctx, dsn)
|
|||
|
|
if err != nil {
|
|||
|
|
t.Skipf("pool: %v", err)
|
|||
|
|
}
|
|||
|
|
if err := pool.Ping(ctx); err != nil {
|
|||
|
|
pool.Close()
|
|||
|
|
t.Skipf("database is not reachable: %v", err)
|
|||
|
|
}
|
|||
|
|
t.Cleanup(pool.Close)
|
|||
|
|
|
|||
|
|
tx, err := pool.Begin(ctx)
|
|||
|
|
if err != nil {
|
|||
|
|
t.Fatalf("begin tx: %v", err)
|
|||
|
|
}
|
|||
|
|
t.Cleanup(func() { _ = tx.Rollback(ctx) })
|
|||
|
|
|
|||
|
|
return tx, ctx
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// requireRateLimits — таблица из миграции 0009. Появляется после первого
|
|||
|
|
// запуска API (db.Migrate на старте) либо go run ./cmd/migrate.
|
|||
|
|
func requireRateLimits(t *testing.T, q Querier, ctx context.Context) {
|
|||
|
|
t.Helper()
|
|||
|
|
var name *string
|
|||
|
|
if err := q.QueryRow(ctx, `SELECT to_regclass('rate_limits')::text`).Scan(&name); err != nil {
|
|||
|
|
t.Fatalf("проверка таблицы rate_limits: %v", err)
|
|||
|
|
}
|
|||
|
|
if name == nil {
|
|||
|
|
t.Fatal("таблицы rate_limits нет — примени миграции: go run ./cmd/migrate")
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// TestDBLimiter — счётчик в БД: потолок держится, ключи независимы,
|
|||
|
|
// окно сбрасывается по времени.
|
|||
|
|
func TestDBLimiter(t *testing.T) {
|
|||
|
|
q, ctx := newTestQuerier(t)
|
|||
|
|
requireRateLimits(t, q, ctx)
|
|||
|
|
|
|||
|
|
l := NewDB(q, "test-login", 2, time.Minute)
|
|||
|
|
ip := uuid.NewString() // уникальный ключ: параллельные тесты не пересекаются
|
|||
|
|
|
|||
|
|
if !l.Allow(ip) {
|
|||
|
|
t.Fatal("1-я попытка: want allowed")
|
|||
|
|
}
|
|||
|
|
if !l.Allow(ip) {
|
|||
|
|
t.Fatal("2-я попытка: want allowed")
|
|||
|
|
}
|
|||
|
|
if l.Allow(ip) {
|
|||
|
|
t.Error("3-я попытка: want denied")
|
|||
|
|
}
|
|||
|
|
// Ключи независимы: чужой бюджет не задет.
|
|||
|
|
if !l.Allow(uuid.NewString()) {
|
|||
|
|
t.Error("другой ключ: want allowed")
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// Окно: кончилось — счётчик начинается заново.
|
|||
|
|
w := NewDB(q, "test-window", 1, 50*time.Millisecond)
|
|||
|
|
k := uuid.NewString()
|
|||
|
|
if !w.Allow(k) {
|
|||
|
|
t.Fatal("1-я попытка в окне: want allowed")
|
|||
|
|
}
|
|||
|
|
if w.Allow(k) {
|
|||
|
|
t.Error("2-я попытка в окне: want denied")
|
|||
|
|
}
|
|||
|
|
time.Sleep(80 * time.Millisecond)
|
|||
|
|
if !w.Allow(k) {
|
|||
|
|
t.Error("после окна: want allowed")
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// TestDBLimiterFailOpen — БД недоступна → запрос разрешён (fail-open)
|
|||
|
|
// и вход не встаёт вместе с базой.
|
|||
|
|
func TestDBLimiterFailOpen(t *testing.T) {
|
|||
|
|
l := NewDB(deadQuerier{}, "test-dead", 1, time.Minute)
|
|||
|
|
if !l.Allow("any") {
|
|||
|
|
t.Error("при ошибке БД лимитер должен разрешить запрос")
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// deadQuerier — «база умерла»: любая операция возвращает ошибку.
|
|||
|
|
type deadQuerier struct{}
|
|||
|
|
|
|||
|
|
func (deadQuerier) Exec(context.Context, string, ...any) (pgconn.CommandTag, error) {
|
|||
|
|
return pgconn.CommandTag{}, context.DeadlineExceeded
|
|||
|
|
}
|
|||
|
|
func (deadQuerier) QueryRow(context.Context, string, ...any) pgx.Row {
|
|||
|
|
return errRow{context.DeadlineExceeded}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// errRow — pgx.Row, у которого всегда ошибка.
|
|||
|
|
type errRow struct{ err error }
|
|||
|
|
|
|||
|
|
func (r errRow) Scan(...any) error { return r.err }
|
|||
|
|
|
|||
|
|
// TestRunCleanup — протухшие строки вычищаются, свежие остаются.
|
|||
|
|
func TestRunCleanup(t *testing.T) {
|
|||
|
|
q, ctx := newTestQuerier(t)
|
|||
|
|
requireRateLimits(t, q, ctx)
|
|||
|
|
|
|||
|
|
fresh := "cleanup-fresh|" + uuid.NewString()
|
|||
|
|
stale := "cleanup-stale|" + uuid.NewString()
|
|||
|
|
// Порог чистки — 25 часов (самое длинное окно лимитов — сутки):
|
|||
|
|
// протухшей считается строка старше порога, свежая — минутная.
|
|||
|
|
if _, err := q.Exec(ctx,
|
|||
|
|
`INSERT INTO rate_limits (key, window_start, count) VALUES ($1, now(), 1), ($2, now() - interval '26 hours', 1)`,
|
|||
|
|
fresh, stale); err != nil {
|
|||
|
|
t.Fatalf("seed: %v", err)
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// Опрос и чистильщик ходят в один tx из разных горутин, а pgx.Tx
|
|||
|
|
// (одно соединение) параллельных запросов не переносит —
|
|||
|
|
// сериализуем мьютексом через адаптер (см. syncQuerier).
|
|||
|
|
sq := &syncQuerier{q: q}
|
|||
|
|
|
|||
|
|
runCtx, cancel := context.WithCancel(context.Background())
|
|||
|
|
defer cancel()
|
|||
|
|
go RunCleanup(runCtx, sq, 20*time.Millisecond)
|
|||
|
|
|
|||
|
|
// Ждём не больше секунды, пока тикер удалит протухшую строку.
|
|||
|
|
deadline := time.Now().Add(time.Second)
|
|||
|
|
for {
|
|||
|
|
var n int
|
|||
|
|
if err := sq.QueryRow(ctx, `SELECT count(*) FROM rate_limits WHERE key = $1`, stale).Scan(&n); err != nil {
|
|||
|
|
t.Fatalf("проверка: %v", err)
|
|||
|
|
}
|
|||
|
|
if n == 0 {
|
|||
|
|
break
|
|||
|
|
}
|
|||
|
|
if time.Now().After(deadline) {
|
|||
|
|
t.Fatal("протухшая строка не вычищена за 1s")
|
|||
|
|
}
|
|||
|
|
time.Sleep(20 * time.Millisecond)
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// Свежая строка не задета.
|
|||
|
|
var n int
|
|||
|
|
if err := sq.QueryRow(ctx, `SELECT count(*) FROM rate_limits WHERE key = $1`, fresh).Scan(&n); err != nil {
|
|||
|
|
t.Fatalf("проверка свежей: %v", err)
|
|||
|
|
}
|
|||
|
|
if n != 1 {
|
|||
|
|
t.Errorf("свежая строка: count = %d, want 1", n)
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// syncQuerier — сериализует запросы к тестовому tx: чистильщик и опрос
|
|||
|
|
// идут из разных горутин, а pgx.Tx параллельных запросов не переносит
|
|||
|
|
// («conn busy»). Мьютекс держится до конца запроса: QueryRow откладывает
|
|||
|
|
// реальный вызов до Scan — только там строка реально выполняется.
|
|||
|
|
type syncQuerier struct {
|
|||
|
|
mu sync.Mutex
|
|||
|
|
q Querier
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func (s *syncQuerier) Exec(ctx context.Context, sql string, args ...any) (pgconn.CommandTag, error) {
|
|||
|
|
s.mu.Lock()
|
|||
|
|
defer s.mu.Unlock()
|
|||
|
|
return s.q.Exec(ctx, sql, args...)
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func (s *syncQuerier) QueryRow(ctx context.Context, sql string, args ...any) pgx.Row {
|
|||
|
|
return &syncRow{s: s, ctx: ctx, sql: sql, args: args}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// syncRow — Scan держит мьютекс адаптера на всё время запроса.
|
|||
|
|
type syncRow struct {
|
|||
|
|
s *syncQuerier
|
|||
|
|
ctx context.Context
|
|||
|
|
sql string
|
|||
|
|
args []any
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func (r *syncRow) Scan(dest ...any) error {
|
|||
|
|
r.s.mu.Lock()
|
|||
|
|
defer r.s.mu.Unlock()
|
|||
|
|
return r.s.q.QueryRow(r.ctx, r.sql, r.args...).Scan(dest...)
|
|||
|
|
}
|