//go:build integration package usage_test import ( "context" "database/sql" "fmt" "sync" "testing" "time" _ "github.com/go-sql-driver/mysql" "github.com/redis/go-redis/v9" "github.com/testcontainers/testcontainers-go" tcmysql "github.com/testcontainers/testcontainers-go/modules/mysql" tcredis "github.com/testcontainers/testcontainers-go/modules/redis" "github.com/wangjia/pangolin/server/internal/usage" ) // -------------------------------------------------------------------------- // Container setup // -------------------------------------------------------------------------- func setupMySQL(t *testing.T) *sql.DB { t.Helper() ctx := context.Background() ctr, err := tcmysql.Run(ctx, "mysql:8.0", tcmysql.WithDatabase("pangolin_test"), tcmysql.WithUsername("root"), tcmysql.WithPassword("test"), ) testcontainers.CleanupContainer(t, ctr) if err != nil { t.Fatalf("mysql container: %v", err) } dsn, err := ctr.ConnectionString(ctx, "parseTime=true", "loc=UTC", "time_zone=%27%2B00%3A00%27") if err != nil { t.Fatalf("mysql dsn: %v", err) } db, err := sql.Open("mysql", dsn) if err != nil { t.Fatalf("open mysql: %v", err) } t.Cleanup(func() { db.Close() }) if err := applySchema(db); err != nil { t.Fatalf("schema: %v", err) } return db } func setupRedis(t *testing.T) *redis.Client { t.Helper() ctx := context.Background() ctr, err := tcredis.Run(ctx, "redis:7-alpine") testcontainers.CleanupContainer(t, ctr) if err != nil { t.Fatalf("redis container: %v", err) } addr, err := ctr.ConnectionString(ctx) if err != nil { t.Fatalf("redis addr: %v", err) } for _, p := range []string{"redis://", "rediss://"} { if len(addr) > len(p) && addr[:len(p)] == p { addr = addr[len(p):] } } rdb := redis.NewClient(&redis.Options{Addr: addr}) t.Cleanup(func() { rdb.Close() }) return rdb } func applySchema(db *sql.DB) error { stmts := []string{ `CREATE TABLE IF NOT EXISTS users ( id BIGINT UNSIGNED AUTO_INCREMENT PRIMARY KEY, uuid CHAR(36) NOT NULL UNIQUE, email VARCHAR(255) NOT NULL UNIQUE, pw_hash VARCHAR(255) NOT NULL DEFAULT '', dp_uuid CHAR(36) NOT NULL, status ENUM('active','banned') NOT NULL DEFAULT 'active', created_at DATETIME(6) NOT NULL DEFAULT CURRENT_TIMESTAMP(6), INDEX idx_dp (dp_uuid) ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4`, `CREATE TABLE IF NOT EXISTS plans ( id BIGINT UNSIGNED AUTO_INCREMENT PRIMARY KEY, code ENUM('free','pro','team') NOT NULL UNIQUE, max_devices INT NOT NULL DEFAULT 1, daily_minutes INT NULL, ad_gate BOOLEAN NOT NULL DEFAULT FALSE ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4`, `CREATE TABLE IF NOT EXISTS subscriptions ( id BIGINT UNSIGNED AUTO_INCREMENT PRIMARY KEY, user_id BIGINT UNSIGNED NOT NULL, plan_id BIGINT UNSIGNED NOT NULL, expires_at DATETIME(6) NOT NULL, source ENUM('trial','code') NOT NULL, created_at DATETIME(6) NOT NULL DEFAULT CURRENT_TIMESTAMP(6), INDEX idx_user_exp (user_id, expires_at) ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4`, `CREATE TABLE IF NOT EXISTS usage_daily ( user_id BIGINT UNSIGNED NOT NULL, date DATE NOT NULL, bytes_up BIGINT UNSIGNED NOT NULL DEFAULT 0, bytes_down BIGINT UNSIGNED NOT NULL DEFAULT 0, minutes_used INT NOT NULL DEFAULT 0, ad_unlocked_at DATETIME(6) NULL, PRIMARY KEY (user_id, date) ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4`, `CREATE TABLE IF NOT EXISTS usage_hourly ( user_id BIGINT UNSIGNED NOT NULL, hour BIGINT NOT NULL, bytes_up BIGINT UNSIGNED NOT NULL DEFAULT 0, bytes_down BIGINT UNSIGNED NOT NULL DEFAULT 0, minutes_used INT NOT NULL DEFAULT 0, PRIMARY KEY (user_id, hour) ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4`, `INSERT IGNORE INTO plans (code, max_devices, daily_minutes, ad_gate) VALUES ('free',1,10,TRUE), ('pro',5,NULL,FALSE), ('team',10,NULL,FALSE)`, } for _, s := range stmts { if _, err := db.Exec(s); err != nil { return fmt.Errorf("exec: %w\nSQL: %s", err, s) } } return nil } // createUser inserts a user with the given dp_uuid and returns its id. func createUser(t *testing.T, db *sql.DB, dpUUID string) int64 { t.Helper() uuid := fmt.Sprintf("u-%d", time.Now().UnixNano()) res, err := db.Exec( `INSERT INTO users (uuid, email, pw_hash, dp_uuid) VALUES (?,?,?,?)`, uuid, uuid+"@example.com", "x", dpUUID) if err != nil { t.Fatalf("createUser: %v", err) } id, _ := res.LastInsertId() return id } // giveSubscription grants userID an active subscription on the given plan. func giveSubscription(t *testing.T, db *sql.DB, userID int64, plan string) { t.Helper() var planID int64 if err := db.QueryRow(`SELECT id FROM plans WHERE code=?`, plan).Scan(&planID); err != nil { t.Fatalf("plan lookup: %v", err) } if _, err := db.Exec( `INSERT INTO subscriptions (user_id, plan_id, expires_at, source) VALUES (?,?,?,'code')`, userID, planID, time.Now().UTC().AddDate(0, 0, 30)); err != nil { t.Fatalf("giveSubscription: %v", err) } } // fakeVerifier is an AdVerifier test double. type fakeVerifier struct{ ok bool } func (f fakeVerifier) Provider() string { return "fake" } func (f fakeVerifier) Verify(context.Context, usage.AdVerifyRequest) error { if f.ok { return nil } return fmt.Errorf("fake verifier: rejected") } // -------------------------------------------------------------------------- // Aggregation // -------------------------------------------------------------------------- func TestIntAggregateConcurrentMultiNode(t *testing.T) { db := setupMySQL(t) rdb := setupRedis(t) store := usage.NewStore(db) agg := usage.NewAggregator(store, rdb, time.Minute, time.Minute) uid := createUser(t, db, "dp-agg-1") const goroutines = 25 const each = 4 var wg sync.WaitGroup wg.Add(goroutines) for g := 0; g < goroutines; g++ { go func(g int) { defer wg.Done() for i := 0; i < each; i++ { err := agg.ReportUsage(context.Background(), usage.Report{ DPUUID: "dp-agg-1", BytesUp: 100, BytesDown: 200, Minutes: 1, BatchID: fmt.Sprintf("b-%d-%d", g, i), // unique → all counted }) if err != nil { t.Errorf("report: %v", err) } } }(g) } wg.Wait() var up, down uint64 var minutes int db.QueryRow(`SELECT bytes_up, bytes_down, minutes_used FROM usage_daily WHERE user_id=?`, uid). Scan(&up, &down, &minutes) wantUp := uint64(goroutines * each * 100) wantMin := goroutines * each if up != wantUp || down != wantUp*2 || minutes != wantMin { t.Errorf("got up=%d down=%d min=%d, want up=%d down=%d min=%d", up, down, minutes, wantUp, wantUp*2, wantMin) } } func TestIntAggregateReplayDedup(t *testing.T) { db := setupMySQL(t) rdb := setupRedis(t) store := usage.NewStore(db) agg := usage.NewAggregator(store, rdb, time.Minute, time.Minute) uid := createUser(t, db, "dp-replay") r := usage.Report{DPUUID: "dp-replay", BytesUp: 500, Minutes: 5, BatchID: "same-batch"} for i := 0; i < 4; i++ { if err := agg.ReportUsage(context.Background(), r); err != nil { t.Fatalf("report %d: %v", i, err) } } var up uint64 var minutes int db.QueryRow(`SELECT bytes_up, minutes_used FROM usage_daily WHERE user_id=?`, uid).Scan(&up, &minutes) if up != 500 || minutes != 5 { t.Errorf("replay not deduped: up=%d minutes=%d, want 500/5", up, minutes) } } func TestIntAggregateCrossUTCDayBucketing(t *testing.T) { db := setupMySQL(t) store := usage.NewStore(db) agg := usage.NewAggregator(store, nil, time.Minute, time.Minute) uid := createUser(t, db, "dp-days") day1 := time.Date(2026, 6, 12, 23, 30, 0, 0, time.UTC) day2 := time.Date(2026, 6, 13, 0, 15, 0, 0, time.UTC) must := func(at time.Time, minutes int) { if err := agg.ReportUsage(context.Background(), usage.Report{ DPUUID: "dp-days", BytesUp: 10, Minutes: minutes, At: at, }); err != nil { t.Fatalf("report: %v", err) } } must(day1, 3) must(day1.Add(10*time.Minute), 2) // still 06-12 must(day2, 7) // 06-13 check := func(date string, wantMin int) { var m int err := db.QueryRow(`SELECT minutes_used FROM usage_daily WHERE user_id=? AND date=?`, uid, date).Scan(&m) if err != nil { t.Fatalf("query %s: %v", date, err) } if m != wantMin { t.Errorf("date %s minutes=%d, want %d", date, m, wantMin) } } check("2026-06-12", 5) check("2026-06-13", 7) } func TestIntAggregateUnknownDPUUIDDropped(t *testing.T) { db := setupMySQL(t) store := usage.NewStore(db) agg := usage.NewAggregator(store, nil, time.Minute, time.Minute) err := agg.ReportUsage(context.Background(), usage.Report{ DPUUID: "does-not-exist", BytesUp: 1, Minutes: 1, }) if err != nil { t.Fatalf("unknown dp_uuid should be a silent no-op, got %v", err) } var n int db.QueryRow(`SELECT COUNT(*) FROM usage_daily`).Scan(&n) if n != 0 { t.Errorf("expected no rows for unknown dp_uuid, got %d", n) } } // -------------------------------------------------------------------------- // Quota (CheckFreeConnect) // -------------------------------------------------------------------------- func TestIntQuotaFreeNotUnlocked(t *testing.T) { db := setupMySQL(t) svc := usage.NewService(usage.NewStore(db), nil, fakeVerifier{ok: true}, time.Hour) uid := createUser(t, db, "dp-q1") // no subscription → free plan _, apiErr := svc.CheckFreeConnect(context.Background(), uid) if apiErr == nil || apiErr.Code != "AD_NOT_UNLOCKED" { t.Fatalf("expected AD_NOT_UNLOCKED, got %v", apiErr) } if apiErr.MessageZH == "" || apiErr.MessageEn == "" { t.Error("expected bilingual error messages") } } func TestIntQuotaFreeUnlockedRemainingDecrements(t *testing.T) { db := setupMySQL(t) store := usage.NewStore(db) agg := usage.NewAggregator(store, nil, time.Minute, time.Minute) svc := usage.NewService(store, nil, fakeVerifier{ok: true}, time.Hour) uid := createUser(t, db, "dp-q2") if _, err := store.MarkAdUnlocked(context.Background(), uid, time.Now().UTC()); err != nil { t.Fatalf("unlock: %v", err) } rem, apiErr := svc.CheckFreeConnect(context.Background(), uid) if apiErr != nil { t.Fatalf("unexpected err: %v", apiErr) } if rem != 10 { t.Errorf("remaining=%d, want 10", rem) } // Consume 4 minutes. agg.ReportUsage(context.Background(), usage.Report{DPUUID: "dp-q2", Minutes: 4}) rem, _ = svc.CheckFreeConnect(context.Background(), uid) if rem != 6 { t.Errorf("after 4 min: remaining=%d, want 6", rem) } } func TestIntQuotaFreeExhausted(t *testing.T) { db := setupMySQL(t) store := usage.NewStore(db) svc := usage.NewService(store, nil, fakeVerifier{ok: true}, time.Hour) uid := createUser(t, db, "dp-q3") store.MarkAdUnlocked(context.Background(), uid, time.Now().UTC()) store.AggregateUsage(context.Background(), uid, time.Now().UTC(), 0, 0, 10) _, apiErr := svc.CheckFreeConnect(context.Background(), uid) if apiErr == nil || apiErr.Code != "QUOTA_EXHAUSTED" { t.Fatalf("expected QUOTA_EXHAUSTED, got %v", apiErr) } } func TestIntQuotaPaidUnlimited(t *testing.T) { db := setupMySQL(t) store := usage.NewStore(db) svc := usage.NewService(store, nil, fakeVerifier{ok: true}, time.Hour) for _, plan := range []string{"pro", "team"} { uid := createUser(t, db, "dp-"+plan) giveSubscription(t, db, uid, plan) // Even with lots of minutes used, paid plans are not gated. store.AggregateUsage(context.Background(), uid, time.Now().UTC(), 0, 0, 9999) rem, apiErr := svc.CheckFreeConnect(context.Background(), uid) if apiErr != nil { t.Fatalf("%s: unexpected err %v", plan, apiErr) } if rem != usage.Unlimited { t.Errorf("%s: remaining=%d, want Unlimited(%d)", plan, rem, usage.Unlimited) } } } // -------------------------------------------------------------------------- // Ads unlock // -------------------------------------------------------------------------- func TestIntAdsUnlockForgedRejected(t *testing.T) { db := setupMySQL(t) rdb := setupRedis(t) svc := usage.NewService(usage.NewStore(db), rdb, fakeVerifier{ok: false}, time.Hour) uid := createUser(t, db, "dp-ad1") _, apiErr := svc.UnlockAd(context.Background(), uid, "device-1", "forged-token") if apiErr == nil || apiErr.Code != "AD_VERIFY_FAILED" { t.Fatalf("expected AD_VERIFY_FAILED, got %v", apiErr) } } func TestIntAdsUnlockReplayRejected(t *testing.T) { db := setupMySQL(t) rdb := setupRedis(t) svc := usage.NewService(usage.NewStore(db), rdb, fakeVerifier{ok: true}, time.Hour) uid := createUser(t, db, "dp-ad2") already, apiErr := svc.UnlockAd(context.Background(), uid, "device-1", "token-xyz") if apiErr != nil || already { t.Fatalf("first unlock: already=%v err=%v", already, apiErr) } // Same token again → replay. _, apiErr = svc.UnlockAd(context.Background(), uid, "device-1", "token-xyz") if apiErr == nil || apiErr.Code != "AD_TOKEN_REPLAY" { t.Fatalf("expected AD_TOKEN_REPLAY, got %v", apiErr) } } func TestIntAdsUnlockSecondTimeIdempotent(t *testing.T) { db := setupMySQL(t) rdb := setupRedis(t) svc := usage.NewService(usage.NewStore(db), rdb, fakeVerifier{ok: true}, time.Hour) uid := createUser(t, db, "dp-ad3") if _, apiErr := svc.UnlockAd(context.Background(), uid, "d", "tok-1"); apiErr != nil { t.Fatalf("first unlock: %v", apiErr) } // Different (valid) token, same day → idempotent success. already, apiErr := svc.UnlockAd(context.Background(), uid, "d", "tok-2") if apiErr != nil { t.Fatalf("second unlock err: %v", apiErr) } if !already { t.Error("expected already-unlocked idempotent success") } // Exactly one unlock timestamp. var n int db.QueryRow(`SELECT COUNT(*) FROM usage_daily WHERE user_id=? AND ad_unlocked_at IS NOT NULL`, uid).Scan(&n) if n != 1 { t.Errorf("expected 1 unlocked day, got %d", n) } } // -------------------------------------------------------------------------- // Usage curve & today summary // -------------------------------------------------------------------------- func TestIntUsageCurveZeroFill(t *testing.T) { db := setupMySQL(t) store := usage.NewStore(db) svc := usage.NewService(store, nil, nil, time.Hour) uid := createUser(t, db, "dp-curve") today := time.Now().UTC() epochHour := func(tm time.Time) int64 { return tm.UTC().Unix() / 3600 } // Only two of the last 7 days have data (curve reads hourly buckets now). db.Exec(`INSERT INTO usage_hourly (user_id, hour, bytes_up, bytes_down, minutes_used) VALUES (?,?,?,?,?)`, uid, epochHour(today), 1000, 2000, 5) db.Exec(`INSERT INTO usage_hourly (user_id, hour, bytes_up, bytes_down, minutes_used) VALUES (?,?,?,?,?)`, uid, epochHour(today.AddDate(0, 0, -3)), 50, 60, 2) points, apiErr := svc.UsageCurve(context.Background(), uid, 7, 0, "") if apiErr != nil { t.Fatalf("curve: %v", apiErr) } if len(points) != 7 { t.Fatalf("expected 7 points, got %d", len(points)) } // Oldest first, dates contiguous. for i := 1; i < len(points); i++ { if points[i].Date <= points[i-1].Date { t.Errorf("points not ascending: %s then %s", points[i-1].Date, points[i].Date) } } // Last point = today with the data. last := points[len(points)-1] if last.BytesUp != 1000 || last.BytesDown != 2000 || last.MinutesUsed != 5 { t.Errorf("today point wrong: %+v", last) } // Days with no data are zero-filled. zeros := 0 for _, p := range points { if p.BytesUp == 0 && p.BytesDown == 0 && p.MinutesUsed == 0 { zeros++ } } if zeros != 5 { t.Errorf("expected 5 zero-filled days, got %d", zeros) } } func TestIntUsageCurveOnlyCurrentUser(t *testing.T) { db := setupMySQL(t) store := usage.NewStore(db) svc := usage.NewService(store, nil, nil, time.Hour) uA := createUser(t, db, "dp-A") uB := createUser(t, db, "dp-B") h := time.Now().UTC().Unix() / 3600 db.Exec(`INSERT INTO usage_hourly (user_id, hour, bytes_up, bytes_down, minutes_used) VALUES (?,?,?,?,?)`, uA, h, 111, 0, 1) db.Exec(`INSERT INTO usage_hourly (user_id, hour, bytes_up, bytes_down, minutes_used) VALUES (?,?,?,?,?)`, uB, h, 999, 0, 9) points, _ := svc.UsageCurve(context.Background(), uA, 1, 0, "") if len(points) != 1 || points[0].BytesUp != 111 { t.Errorf("user A curve leaked other user data: %+v", points) } } func TestIntTodaySummary(t *testing.T) { db := setupMySQL(t) store := usage.NewStore(db) svc := usage.NewService(store, nil, nil, time.Hour) // Free user: limited + ad flag. uFree := createUser(t, db, "dp-sf") store.MarkAdUnlocked(context.Background(), uFree, time.Now().UTC()) store.AggregateUsage(context.Background(), uFree, time.Now().UTC(), 0, 0, 3) sum, apiErr := svc.TodaySummary(context.Background(), uFree) if apiErr != nil { t.Fatalf("summary: %v", apiErr) } if sum.MinutesUsed != 3 || sum.MinutesRemaining == nil || *sum.MinutesRemaining != 7 || !sum.AdUnlocked { t.Errorf("free summary wrong: %+v (remaining=%v)", sum, sum.MinutesRemaining) } // Pro user: unlimited (nil remaining). uPro := createUser(t, db, "dp-sp") giveSubscription(t, db, uPro, "pro") sum2, _ := svc.TodaySummary(context.Background(), uPro) if sum2.MinutesRemaining != nil { t.Errorf("pro remaining should be nil(unlimited), got %v", *sum2.MinutesRemaining) } }