new-api/service/rebate_test.go

package service

import (
	"fmt"
	"strings"
	"testing"
	"time"

	"github.com/QuantumNous/new-api/common"
	"github.com/QuantumNous/new-api/model"
	"github.com/glebarez/sqlite"
	"github.com/stretchr/testify/require"
	"gorm.io/gorm"
)

func setupRebateTestDB(t *testing.T) *gorm.DB {
	t.Helper()

	previousDB := model.DB
	previousLogDB := model.LOG_DB
	common.OptionMapRWMutex.RLock()
	previousOptionMap := common.OptionMap
	common.OptionMapRWMutex.RUnlock()

	common.UsingSQLite = true
	common.UsingMySQL = false
	common.UsingPostgreSQL = false
	common.RedisEnabled = false

	dsn := fmt.Sprintf("file:%s?mode=memory&cache=shared", strings.ReplaceAll(t.Name(), "/", "_"))
	db, err := gorm.Open(sqlite.Open(dsn), &gorm.Config{})
	require.NoError(t, err)
	require.NoError(t, db.AutoMigrate(&model.User{}, &model.RebateRecord{}, &model.Log{}))

	model.DB = db
	model.LOG_DB = db
	setRebateOptions(false, "0", "0")

	t.Cleanup(func() {
		model.DB = previousDB
		model.LOG_DB = previousLogDB
		common.OptionMapRWMutex.Lock()
		common.OptionMap = previousOptionMap
		common.OptionMapRWMutex.Unlock()
		sqlDB, err := db.DB()
		if err == nil {
			_ = sqlDB.Close()
		}
	})

	return db
}

func setRebateOptions(enabled bool, ratePercent string, frozenDays string) {
	common.OptionMapRWMutex.Lock()
	defer common.OptionMapRWMutex.Unlock()
	common.OptionMap = map[string]string{
		model.AffRebateEnabledKey:     fmt.Sprintf("%t", enabled),
		model.AffRebateRatePercentKey: ratePercent,
		model.AffRebateFrozenDaysKey:  frozenDays,
	}
}

func createRebateUsers(t *testing.T, db *gorm.DB, inviterOverrides ...func(*model.User)) (*model.User, *model.User) {
	t.Helper()

	inviter := &model.User{
		Id:       1,
		Username: "inviter_" + strings.ReplaceAll(t.Name(), "/", "_"),
		Password: "password123",
		AffCode:  "aff_inviter_" + strings.ReplaceAll(t.Name(), "/", "_"),
		Status:   common.UserStatusEnabled,
	}
	for _, override := range inviterOverrides {
		override(inviter)
	}
	invitee := &model.User{
		Id:        2,
		Username:  "invitee_" + strings.ReplaceAll(t.Name(), "/", "_"),
		Password:  "password123",
		AffCode:   "aff_invitee_" + strings.ReplaceAll(t.Name(), "/", "_"),
		Status:    common.UserStatusEnabled,
		InviterId: inviter.Id,
	}
	require.NoError(t, db.Create(inviter).Error)
	require.NoError(t, db.Create(invitee).Error)
	return inviter, invitee
}

func loadRebateUser(t *testing.T, db *gorm.DB, id int) model.User {
	t.Helper()
	var user model.User
	require.NoError(t, db.First(&user, id).Error)
	return user
}

func TestProcessRebateAfterRechargeSkipsWhenDisabled(t *testing.T) {
	db := setupRebateTestDB(t)
	_, invitee := createRebateUsers(t, db)
	setRebateOptions(false, "10", "7")

	require.NoError(t, ProcessRebateAfterRecharge(invitee, 1001, "topup", 1000))

	var count int64
	require.NoError(t, db.Model(&model.RebateRecord{}).Count(&count).Error)
	require.Equal(t, int64(0), count)
	inviter := loadRebateUser(t, db, 1)
	require.Equal(t, 0, inviter.AffQuota)
	require.Equal(t, 0, inviter.AffHistoryQuota)
}

func TestProcessRebateAfterRechargeRejectsMissingOrderID(t *testing.T) {
	db := setupRebateTestDB(t)
	_, invitee := createRebateUsers(t, db)
	setRebateOptions(true, "10", "7")

	err := ProcessRebateAfterRecharge(invitee, 0, "topup", 1000)

	require.Error(t, err)
	require.Contains(t, err.Error(), "order id")
	var count int64
	require.NoError(t, db.Model(&model.RebateRecord{}).Count(&count).Error)
	require.Equal(t, int64(0), count)
}

func TestProcessRebateAfterRechargeUsesGlobalRate(t *testing.T) {
	db := setupRebateTestDB(t)
	_, invitee := createRebateUsers(t, db)
	setRebateOptions(true, "10", "7")

	require.NoError(t, ProcessRebateAfterRecharge(invitee, 1002, "topup", 1000))

	var record model.RebateRecord
	require.NoError(t, db.First(&record).Error)
	require.Equal(t, 1, record.InviterId)
	require.Equal(t, 2, record.InviteeId)
	require.Equal(t, 1002, record.OrderId)
	require.Equal(t, "topup", record.OrderType)
	require.Equal(t, 1000, record.OrderAmount)
	require.Equal(t, 100, record.RebateAmount)
	require.Equal(t, 10, record.RatePercent)
	require.Equal(t, "frozen", record.Status)
	require.NotNil(t, record.FrozenUntil)
	require.Nil(t, record.ReleasedAt)

	inviter := loadRebateUser(t, db, 1)
	require.Equal(t, 0, inviter.AffQuota)
	require.Equal(t, 100, inviter.AffHistoryQuota)
}

func TestProcessRebateAfterRechargeUsesPerUserOverrides(t *testing.T) {
	db := setupRebateTestDB(t)
	_, invitee := createRebateUsers(t, db, func(inviter *model.User) {
		inviter.AffRebateRatePercent = 25
		inviter.AffRebateFrozenDays = 1
	})
	setRebateOptions(true, "10", "7")

	require.NoError(t, ProcessRebateAfterRecharge(invitee, 1003, "topup", 1000))

	var record model.RebateRecord
	require.NoError(t, db.First(&record).Error)
	require.Equal(t, 250, record.RebateAmount)
	require.Equal(t, 25, record.RatePercent)
	require.Equal(t, "frozen", record.Status)
	require.NotNil(t, record.FrozenUntil)
	require.WithinDuration(t, time.Now().Add(24*time.Hour), *record.FrozenUntil, time.Minute)

	inviter := loadRebateUser(t, db, 1)
	require.Equal(t, 0, inviter.AffQuota)
	require.Equal(t, 250, inviter.AffHistoryQuota)
}

func TestProcessRebateAfterRechargeIsIdempotentForSameOrderAndUsers(t *testing.T) {
	db := setupRebateTestDB(t)
	_, invitee := createRebateUsers(t, db)
	setRebateOptions(true, "10", "0")

	require.NoError(t, ProcessRebateAfterRecharge(invitee, 1004, "topup", 1000))
	require.NoError(t, ProcessRebateAfterRecharge(invitee, 1004, "topup", 1000))

	var count int64
	require.NoError(t, db.Model(&model.RebateRecord{}).Count(&count).Error)
	require.Equal(t, int64(1), count)
	inviter := loadRebateUser(t, db, 1)
	require.Equal(t, 100, inviter.AffQuota)
	require.Equal(t, 100, inviter.AffHistoryQuota)
}

func TestReleaseExpiredRebatesMarksOnlyDueFrozenRecords(t *testing.T) {
	db := setupRebateTestDB(t)
	inviter, invitee := createRebateUsers(t, db)
	past := time.Now().Add(-time.Hour)
	future := time.Now().Add(time.Hour)
	require.NoError(t, db.Model(inviter).Updates(map[string]interface{}{
		"aff_history": 50,
	}).Error)
	require.NoError(t, db.Create(&model.RebateRecord{
		InviterId:    inviter.Id,
		InviteeId:    invitee.Id,
		OrderId:      1005,
		OrderType:    "topup",
		OrderAmount:  500,
		RebateAmount: 50,
		RatePercent:  10,
		Status:       "frozen",
		FrozenUntil:  &past,
	}).Error)
	require.NoError(t, db.Create(&model.RebateRecord{
		InviterId:    inviter.Id,
		InviteeId:    invitee.Id,
		OrderId:      1006,
		OrderType:    "topup",
		OrderAmount:  500,
		RebateAmount: 50,
		RatePercent:  10,
		Status:       "frozen",
		FrozenUntil:  &future,
	}).Error)

	require.NoError(t, ReleaseExpiredRebates())

	var due model.RebateRecord
	require.NoError(t, db.Where("order_id = ?", 1005).First(&due).Error)
	require.Equal(t, "released", due.Status)
	require.NotNil(t, due.ReleasedAt)

	var notDue model.RebateRecord
	require.NoError(t, db.Where("order_id = ?", 1006).First(&notDue).Error)
	require.Equal(t, "frozen", notDue.Status)
	require.Nil(t, notDue.ReleasedAt)

	reloadedInviter := loadRebateUser(t, db, inviter.Id)
	require.Equal(t, 50, reloadedInviter.AffQuota)
	require.Equal(t, 50, reloadedInviter.AffHistoryQuota)
}

func TestRunRebateReleaseOnceReleasesDueFrozenRecords(t *testing.T) {
	db := setupRebateTestDB(t)
	inviter, invitee := createRebateUsers(t, db)
	past := time.Now().Add(-time.Hour)
	require.NoError(t, db.Create(&model.RebateRecord{
		InviterId:    inviter.Id,
		InviteeId:    invitee.Id,
		OrderId:      1007,
		OrderType:    "topup",
		OrderAmount:  500,
		RebateAmount: 50,
		RatePercent:  10,
		Status:       "frozen",
		FrozenUntil:  &past,
	}).Error)

	require.NoError(t, RunRebateReleaseOnce())

	var record model.RebateRecord
	require.NoError(t, db.Where("order_id = ?", 1007).First(&record).Error)
	require.Equal(t, "released", record.Status)
	require.NotNil(t, record.ReleasedAt)

	reloadedInviter := loadRebateUser(t, db, inviter.Id)
	require.Equal(t, 50, reloadedInviter.AffQuota)
}