hi-server/scripts/test_device_trial_registration.go

//go:build ignore

package main

import (
	"context"
	"flag"
	"fmt"
	"os"
	"strings"
	"time"

	"github.com/perfect-panel/server/initialize"
	"github.com/perfect-panel/server/internal/config"
	authlogic "github.com/perfect-panel/server/internal/logic/auth"
	modelAuth "github.com/perfect-panel/server/internal/model/auth"
	modelLog "github.com/perfect-panel/server/internal/model/log"
	modelNode "github.com/perfect-panel/server/internal/model/node"
	modelSubscribe "github.com/perfect-panel/server/internal/model/subscribe"
	modelSystem "github.com/perfect-panel/server/internal/model/system"
	modelUser "github.com/perfect-panel/server/internal/model/user"
	"github.com/perfect-panel/server/internal/svc"
	"github.com/perfect-panel/server/internal/types"
	"github.com/perfect-panel/server/pkg/conf"
	"github.com/perfect-panel/server/pkg/orm"
	"github.com/perfect-panel/server/pkg/tool"
	"github.com/redis/go-redis/v9"
	"gorm.io/gorm"
)

func main() {
	var (
		configPath = flag.String("config", "etc/ppanel.yaml", "config file path on the test server")
		dsn        = flag.String("dsn", "", "optional MySQL DSN override")
		identifier = flag.String("identifier", "", "optional device identifier; defaults to a unique test identifier")
		ip         = flag.String("ip", "", "optional request IP; defaults to a reserved test IP")
		userAgent  = flag.String("user-agent", "CodexDeviceTrialTest/1.0", "device user agent")
		write      = flag.Bool("write", false, "actually create a test device user by running DeviceLogin")
		cleanup    = flag.Bool("cleanup", false, "delete the test user/device/subscription/log rows after verification")
	)
	flag.Parse()

	if !*write {
		fmt.Println("Refusing to write DB without -write.")
		fmt.Println("Example:")
		fmt.Printf("  go run scripts/test_device_trial_registration.go -config %s -write\n", *configPath)
		os.Exit(2)
	}

	ctx := context.Background()
	cfg := loadConfig(*configPath, *dsn)
	env := mustNewDeviceTrialEnv(ctx, cfg)
	defer env.close()

	initialize.Device(env.svcCtx)
	initialize.Register(env.svcCtx)

	if *identifier == "" {
		*identifier = fmt.Sprintf("codex-device-trial-%d", time.Now().UnixNano())
	}
	if *ip == "" {
		now := time.Now().UnixNano()
		*ip = fmt.Sprintf("198.18.%d.%d", now%200+1, now/200%200+1)
	}

	fmt.Println("== device registration no-trial test ==")
	fmt.Printf("mysql: %s/%s\n", env.cfg.MySQL.Addr, env.cfg.MySQL.Dbname)
	fmt.Printf("redis: %s db=%d\n", env.cfg.Redis.Host, env.cfg.Redis.DB)
	fmt.Printf("device.enable=%v\n", env.svcCtx.Config.Device.Enable)
	fmt.Printf("register.enable_trial=%v trial_subscribe=%d trial_time=%d trial_time_unit=%s\n",
		env.svcCtx.Config.Register.EnableTrial,
		env.svcCtx.Config.Register.TrialSubscribe,
		env.svcCtx.Config.Register.TrialTime,
		env.svcCtx.Config.Register.TrialTimeUnit,
	)
	fmt.Printf("identifier=%s ip=%s user_agent=%s\n", *identifier, *ip, *userAgent)

	if err := ensureIdentifierUnused(ctx, env.db, *identifier); err != nil {
		fail(err)
	}

	logic := authlogic.NewDeviceLoginLogic(ctx, env.svcCtx)
	resp, err := logic.DeviceLogin(&types.DeviceLoginRequest{
		Identifier: *identifier,
		IP:         *ip,
		UserAgent:  *userAgent,
	})
	if err != nil {
		fail(fmt.Errorf("DeviceLogin failed: %w", err))
	}
	if resp == nil || strings.TrimSpace(resp.Token) == "" {
		fail(fmt.Errorf("DeviceLogin returned empty token"))
	}
	fmt.Printf("login token: ok len=%d\n", len(resp.Token))

	device, err := env.svcCtx.UserModel.FindOneDeviceByIdentifier(ctx, *identifier)
	if err != nil {
		fail(fmt.Errorf("query created device failed: %w", err))
	}
	fmt.Printf("device: id=%d sn=%s user_id=%d created_at=%s\n",
		device.Id,
		tool.DeviceIdToHash(device.Id),
		device.UserId,
		device.CreatedAt.Format(time.RFC3339),
	)

	var subs []modelUser.Subscribe
	if err = env.db.WithContext(ctx).
		Where("user_id = ?", device.UserId).
		Order("id ASC").
		Find(&subs).Error; err != nil {
		fail(fmt.Errorf("query user_subscribe failed: %w", err))
	}
	for i := range subs {
		sub := &subs[i]
		fmt.Printf("subscribe: id=%d order_id=%d subscribe_id=%d status=%d start=%s expire=%s token_empty=%v\n",
			sub.Id,
			sub.OrderId,
			sub.SubscribeId,
			sub.Status,
			sub.StartTime.Format(time.RFC3339),
			sub.ExpireTime.Format(time.RFC3339),
			sub.Token == "",
		)
		if sub.OrderId == 0 &&
			sub.SubscribeId == env.svcCtx.Config.Register.TrialSubscribe &&
			(sub.Status == 0 || sub.Status == 1) &&
			sub.ExpireTime.After(time.Now()) {
			fail(fmt.Errorf("FAIL: device registration unexpectedly granted trial user_subscribe_id=%d user_id=%d", sub.Id, device.UserId))
		}
	}

	fmt.Printf("PASS: device registration created no active trial subscription for user_id=%d\n", device.UserId)

	if *cleanup {
		if err = cleanupTestRows(ctx, env.db, device.UserId); err != nil {
			fail(fmt.Errorf("cleanup failed: %w", err))
		}
		fmt.Printf("cleanup: deleted test rows for user_id=%d\n", device.UserId)
	}
}

type deviceTrialEnv struct {
	db     *gorm.DB
	rds    *redis.Client
	cfg    config.Config
	svcCtx *svc.ServiceContext
}

func mustNewDeviceTrialEnv(ctx context.Context, cfg config.Config) *deviceTrialEnv {
	db, err := orm.ConnectMysql(orm.Mysql{Config: cfg.MySQL})
	must(err)

	rds := redis.NewClient(&redis.Options{
		Addr:         cfg.Redis.Host,
		Password:     cfg.Redis.Pass,
		DB:           cfg.Redis.DB,
		PoolSize:     cfg.Redis.PoolSize,
		MinIdleConns: cfg.Redis.MinIdleConns,
	})
	must(rds.Ping(ctx).Err())

	svcCtx := &svc.ServiceContext{
		DB:             db,
		Redis:          rds,
		Config:         cfg,
		AuthModel:      modelAuth.NewModel(db, rds),
		LogModel:       modelLog.NewModel(db),
		NodeModel:      modelNode.NewModel(db, rds),
		SystemModel:    modelSystem.NewModel(db, rds),
		UserModel:      modelUser.NewModel(db, rds),
		SubscribeModel: modelSubscribe.NewModel(db, rds),
	}
	return &deviceTrialEnv{db: db, rds: rds, cfg: cfg, svcCtx: svcCtx}
}

func (e *deviceTrialEnv) close() {
	if e == nil || e.rds == nil {
		return
	}
	_ = e.rds.Close()
}

func loadConfig(path, dsn string) config.Config {
	var cfg config.Config
	conf.MustLoad(path, &cfg)
	if dsn != "" {
		parsed := orm.ParseDSN(dsn)
		if parsed == nil {
			fail(fmt.Errorf("invalid dsn"))
		}
		cfg.MySQL = *parsed
	}
	return cfg
}

func ensureIdentifierUnused(ctx context.Context, db *gorm.DB, identifier string) error {
	var count int64
	if err := db.WithContext(ctx).
		Model(&modelUser.Device{}).
		Where("identifier = ?", identifier).
		Count(&count).Error; err != nil {
		return err
	}
	if count > 0 {
		return fmt.Errorf("identifier already exists: %s", identifier)
	}
	return nil
}

func cleanupTestRows(ctx context.Context, db *gorm.DB, userID int64) error {
	return db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
		if err := tx.Where("object_id = ?", userID).Delete(&modelLog.SystemLog{}).Error; err != nil {
			return err
		}
		if err := tx.Where("user_id = ?", userID).Delete(&modelUser.Subscribe{}).Error; err != nil {
			return err
		}
		if err := tx.Where("user_id = ?", userID).Delete(&modelUser.AuthMethods{}).Error; err != nil {
			return err
		}
		if err := tx.Where("user_id = ?", userID).Delete(&modelUser.Device{}).Error; err != nil {
			return err
		}
		return tx.Where("id = ?", userID).Delete(&modelUser.User{}).Error
	})
}

func must(err error) {
	if err != nil {
		fail(err)
	}
}

func fail(err error) {
	fmt.Fprintln(os.Stderr, err)
	os.Exit(1)
}