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