pangolin/server/internal/store/sqlite_stores_test.go

package store_test

import (
	"context"
	"database/sql"
	"testing"
	"time"

	"github.com/wangjia/pangolin/server/internal/codes"
	"github.com/wangjia/pangolin/server/internal/config"
	dbx "github.com/wangjia/pangolin/server/internal/db"
	"github.com/wangjia/pangolin/server/internal/nodes"
	agentv1 "github.com/wangjia/pangolin/server/internal/pb/agentv1"
	"github.com/wangjia/pangolin/server/internal/provision"
	"github.com/wangjia/pangolin/server/internal/sessions"
	"github.com/wangjia/pangolin/server/internal/store"
	"github.com/wangjia/pangolin/server/internal/usage"
)

// openSQLite returns a freshly-migrated in-memory SQLite DB.
func openSQLite(t *testing.T) *sql.DB {
	t.Helper()
	db, err := store.Open(&config.Config{Driver: "sqlite", DSN: ":memory:"})
	if err != nil {
		t.Fatalf("open: %v", err)
	}
	if err := store.MigrateUp(db, "sqlite"); err != nil {
		t.Fatalf("migrate: %v", err)
	}
	t.Cleanup(func() { _ = db.Close() })
	return db
}

// These tests exercise the dialect-sensitive store methods (upserts, row locks,
// Go-side date math) against real SQLite — the behavioral proof that P2/P3 work.

func TestSQLite_UsageAccumulate(t *testing.T) {
	ctx := context.Background()
	db := openSQLite(t)
	us := usage.NewStore(db)
	day := time.Date(2026, 6, 17, 0, 0, 0, 0, time.UTC)

	if err := us.AggregateUsage(ctx, 1, day, 100, 200, 5); err != nil {
		t.Fatalf("aggregate 1: %v", err)
	}
	if err := us.AggregateUsage(ctx, 1, day, 50, 25, 3); err != nil {
		t.Fatalf("aggregate 2: %v", err)
	}

	var up, down, mins int64
	if err := db.QueryRow(
		`SELECT bytes_up, bytes_down, minutes_used FROM usage_daily WHERE user_id=1`,
	).Scan(&up, &down, &mins); err != nil {
		t.Fatalf("read: %v", err)
	}
	if up != 150 || down != 225 || mins != 8 {
		t.Errorf("accumulate upsert wrong: up=%d down=%d mins=%d, want 150/225/8", up, down, mins)
	}
}

// AddAdBonusMinutes 是免费版累加式看广告加时的核心原语:每次广告 +N 分钟,
// 封顶 ceiling,返回新总额与本次实际加时(封顶后为 0)。
func TestSQLite_AddAdBonusMinutes(t *testing.T) {
	ctx := context.Background()
	db := openSQLite(t)
	us := usage.NewStore(db)
	day := time.Date(2026, 6, 30, 0, 0, 0, 0, time.UTC)

	// First ad on a fresh day → insert row, bonus 10.
	newBonus, granted, err := us.AddAdBonusMinutes(ctx, 7, day, 10, 25)
	if err != nil || newBonus != 10 || granted != 10 {
		t.Fatalf("first: newBonus=%d granted=%d err=%v, want 10/10", newBonus, granted, err)
	}

	// Second ad → accumulate to 20.
	newBonus, granted, err = us.AddAdBonusMinutes(ctx, 7, day, 10, 25)
	if err != nil || newBonus != 20 || granted != 10 {
		t.Fatalf("second: newBonus=%d granted=%d err=%v, want 20/10", newBonus, granted, err)
	}

	// Third ad → clamped at ceiling 25, so only +5 granted.
	newBonus, granted, err = us.AddAdBonusMinutes(ctx, 7, day, 10, 25)
	if err != nil || newBonus != 25 || granted != 5 {
		t.Fatalf("third: newBonus=%d granted=%d err=%v, want 25/5", newBonus, granted, err)
	}

	// Fourth ad → already at ceiling, 0 granted.
	newBonus, granted, err = us.AddAdBonusMinutes(ctx, 7, day, 10, 25)
	if err != nil || newBonus != 25 || granted != 0 {
		t.Fatalf("fourth: newBonus=%d granted=%d err=%v, want 25/0", newBonus, granted, err)
	}

	// Existing usage_daily columns are preserved alongside the bonus.
	if err := us.AggregateUsage(ctx, 7, day, 0, 0, 4); err != nil {
		t.Fatalf("aggregate: %v", err)
	}
	d, err := us.GetDay(ctx, 7, day)
	if err != nil || d == nil {
		t.Fatalf("getday: %v", err)
	}
	if d.AdBonusMinutes != 25 || d.MinutesUsed != 4 {
		t.Errorf("getday wrong: bonus=%d used=%d, want 25/4", d.AdBonusMinutes, d.MinutesUsed)
	}
}

// HasActiveSession 是「近实时远程下线」的服务端判据:有非吊销会话=在线,
// 强制退出(RevokeByDevice)后=离线 → 客户端轮询到即登出。
func TestSQLite_SessionHasActiveSession(t *testing.T) {
	ctx := context.Background()
	db := openSQLite(t)
	// FK 前置:user + device。
	res, err := db.ExecContext(ctx,
		`INSERT INTO users (uuid, email, pw_hash, dp_uuid) VALUES ('u-1','a@b.c','h','dp-1')`)
	if err != nil {
		t.Fatalf("user: %v", err)
	}
	uid, _ := res.LastInsertId()
	res, err = db.ExecContext(ctx,
		`INSERT INTO devices (uuid, user_id, name, platform) VALUES ('d-1', ?, 'Mac', 'macos')`, uid)
	if err != nil {
		t.Fatalf("device: %v", err)
	}
	did, _ := res.LastInsertId()

	ss := sessions.NewStore(db)
	if ok, err := ss.HasActiveSession(ctx, uid, did); err != nil || ok {
		t.Fatalf("无会话应为非活跃: ok=%v err=%v", ok, err)
	}
	if err := ss.Create(ctx, uid, did, "jti-1", "1.2.3.4", "v1.0"); err != nil {
		t.Fatalf("create: %v", err)
	}
	if ok, err := ss.HasActiveSession(ctx, uid, did); err != nil || !ok {
		t.Fatalf("建会话后应活跃: ok=%v err=%v", ok, err)
	}
	if set, err := ss.ActiveSessionDeviceIDs(ctx, uid); err != nil || !set[did] {
		t.Fatalf("活跃设备集应含 %d: set=%v err=%v", did, set, err)
	}
	if _, err := ss.RevokeByDevice(ctx, uid, did); err != nil {
		t.Fatalf("revoke: %v", err)
	}
	if ok, err := ss.HasActiveSession(ctx, uid, did); err != nil || ok {
		t.Fatalf("强制退出后应非活跃: ok=%v err=%v", ok, err)
	}
	if set, err := ss.ActiveSessionDeviceIDs(ctx, uid); err != nil || set[did] {
		t.Fatalf("强退后活跃设备集不应含 %d(→立即离线): set=%v err=%v", did, set, err)
	}
}

func TestSQLite_NodeAccumulateUsage(t *testing.T) {
	ctx := context.Background()
	db := openSQLite(t)
	ns := nodes.NewSQLNodeStore(db)
	day := time.Date(2026, 6, 17, 0, 0, 0, 0, time.UTC)

	if err := ns.AccumulateUsage(ctx, 7, day, 10, 20, 1); err != nil {
		t.Fatalf("accumulate 1: %v", err)
	}
	if err := ns.AccumulateUsage(ctx, 7, day, 5, 5, 2); err != nil {
		t.Fatalf("accumulate 2: %v", err)
	}
	var up, down, mins int64
	if err := db.QueryRow(
		`SELECT bytes_up, bytes_down, minutes_used FROM usage_daily WHERE user_id=7`,
	).Scan(&up, &down, &mins); err != nil {
		t.Fatalf("read: %v", err)
	}
	if up != 15 || down != 25 || mins != 3 {
		t.Errorf("node accumulate wrong: up=%d down=%d mins=%d, want 15/25/3", up, down, mins)
	}
}

func TestSQLite_DirectoryVersionBump(t *testing.T) {
	ctx := context.Background()
	db := openSQLite(t)
	d := dbx.DialectForDB(db)

	// Migration 7 seeds (id=1, version=1). Bump twice → 3.
	for i := 0; i < 2; i++ {
		if err := store.BumpDirectoryVersion(ctx, db, d); err != nil {
			t.Fatalf("bump %d: %v", i, err)
		}
	}
	var v int64
	if err := db.QueryRow(`SELECT version FROM directory_version WHERE id=1`).Scan(&v); err != nil {
		t.Fatalf("read: %v", err)
	}
	if v != 3 {
		t.Errorf("version = %d, want 3", v)
	}
}

func TestSQLite_IdempotencyNoopUpsert(t *testing.T) {
	ctx := context.Background()
	db := openSQLite(t)
	ps := provision.NewMySQLStore(db)

	if err := ps.SaveIdempotency(ctx, "k1", "uuid-a"); err != nil {
		t.Fatalf("save 1: %v", err)
	}
	// Second save with same key must be a no-op (insert-or-ignore), keeping uuid-a.
	if err := ps.SaveIdempotency(ctx, "k1", "uuid-b"); err != nil {
		t.Fatalf("save 2: %v", err)
	}
	got, ok, err := ps.LookupIdempotency(ctx, "k1")
	if err != nil {
		t.Fatalf("lookup: %v", err)
	}
	if !ok || got != "uuid-a" {
		t.Errorf("idempotency no-op failed: got=%q ok=%v, want uuid-a", got, ok)
	}
}

func TestSQLite_PersistCredentialUpsert(t *testing.T) {
	ctx := context.Background()
	db := openSQLite(t)
	seedNode(t, db, 1)
	ns := nodes.NewSQLNodeStore(db)

	exp1 := time.Date(2026, 7, 1, 0, 0, 0, 0, time.UTC)
	cred := &agentv1.Credential{DpUUID: "dp-1", Protocol: 3, Flow: "xtls-rprx-vision"}
	if err := ns.PersistCredential(ctx, 1, cred, exp1); err != nil {
		t.Fatalf("persist 1: %v", err)
	}
	// Upsert same (node_id, dp_uuid) with a later expiry → row updated, not duplicated.
	exp2 := exp1.Add(24 * time.Hour)
	if err := ns.PersistCredential(ctx, 1, cred, exp2); err != nil {
		t.Fatalf("persist 2: %v", err)
	}
	var n int
	if err := db.QueryRow(`SELECT COUNT(*) FROM connect_credentials WHERE node_id=1 AND dp_uuid='dp-1'`).Scan(&n); err != nil {
		t.Fatalf("count: %v", err)
	}
	if n != 1 {
		t.Errorf("credential rows = %d, want 1 (upsert, not insert)", n)
	}
}

func TestSQLite_CodesRedeemFlow(t *testing.T) {
	ctx := context.Background()
	db := openSQLite(t)
	seedUser(t, db, 1)
	cs := codes.NewStore(db)

	// Seed a batch + an 'unused' code for plan 'pro' (seeded id resolved via store).
	planID, err := cs.GetPlanID(ctx, codes.PlanPro)
	if err != nil {
		t.Fatalf("plan id: %v", err)
	}
	batchID, err := cs.CreateBatch(ctx, codes.ChannelManual, "tester", "")
	if err != nil {
		t.Fatalf("create batch: %v", err)
	}
	const hash = "abc123hash"
	if err := cs.CreateCode(ctx, hash, planID, 30, batchID); err != nil {
		t.Fatalf("create code: %v", err)
	}

	// Redeem inside a tx: lock the code (FOR UPDATE on mysql / BEGIN IMMEDIATE on
	// sqlite), mark redeemed, create a subscription with Go-computed expiry.
	tx, err := cs.BeginTx(ctx)
	if err != nil {
		t.Fatalf("begin: %v", err)
	}
	cr, err := cs.FindCodeByHashForUpdate(ctx, tx, hash)
	if err != nil || cr == nil {
		_ = tx.Rollback()
		t.Fatalf("find for update: cr=%v err=%v", cr, err)
	}
	if cr.Status != "unused" {
		_ = tx.Rollback()
		t.Fatalf("code status = %q, want unused", cr.Status)
	}
	if err := cs.MarkRedeemed(ctx, tx, cr.ID, 1); err != nil {
		_ = tx.Rollback()
		t.Fatalf("mark redeemed: %v", err)
	}
	subID, err := cs.CreateSubscription(ctx, tx, 1, planID, 30, time.Time{})
	if err != nil {
		_ = tx.Rollback()
		t.Fatalf("create sub: %v", err)
	}
	// Extend by 10 more days (Go-side max(expires,now)+days).
	if err := cs.ExtendSubscription(ctx, tx, subID, 10); err != nil {
		_ = tx.Rollback()
		t.Fatalf("extend: %v", err)
	}
	if err := tx.Commit(); err != nil {
		t.Fatalf("commit: %v", err)
	}

	// Verify: code redeemed; subscription expires ~40 days out.
	var status string
	if err := db.QueryRow(`SELECT status FROM codes WHERE id=?`, cr.ID).Scan(&status); err != nil {
		t.Fatalf("read code: %v", err)
	}
	if status != "redeemed" {
		t.Errorf("code status = %q, want redeemed", status)
	}
	var expires time.Time
	if err := db.QueryRow(`SELECT expires_at FROM subscriptions WHERE id=?`, subID).Scan(&expires); err != nil {
		t.Fatalf("read sub: %v", err)
	}
	wantMin := time.Now().UTC().AddDate(0, 0, 39)
	if expires.Before(wantMin) {
		t.Errorf("subscription expires_at = %v, want ≥ ~40 days out (%v)", expires, wantMin)
	}
}

// --- seed helpers (raw SQL, satisfy foreign keys) ---

func seedUser(t *testing.T, db *sql.DB, id int64) {
	t.Helper()
	if _, err := db.Exec(
		`INSERT INTO users (id, uuid, email, pw_hash, dp_uuid, status, created_at)
		 VALUES (?, ?, ?, 'x', ?, 'active', ?)`,
		id, "u-uuid", "u@example.com", "dp-u", time.Now().UTC()); err != nil {
		t.Fatalf("seed user: %v", err)
	}
}

func seedNode(t *testing.T, db *sql.DB, id int64) {
	t.Helper()
	if _, err := db.Exec(
		`INSERT INTO providers (id, name, api_kind, regions, pool, enabled)
		 VALUES (1, 'p', 'fake', '[]', 'consumable', 1)`); err != nil {
		t.Fatalf("seed provider: %v", err)
	}
	if _, err := db.Exec(
		`INSERT INTO nodes (id, uuid, region, name_zh, name_en, role, tier, endpoint,
		   reality_pbk, reality_sni, provider_id, status, weight, created_at)
		 VALUES (?, 'n-uuid', 'HK', 'zh', 'en', 'entry', 'pro', '1.2.3.4:443',
		   'pbk', 'www.apple.com', 1, 'up', 100, ?)`,
		id, time.Now().UTC()); err != nil {
		t.Fatalf("seed node: %v", err)
	}
}