hi-server/internal/logic/admin/invite/benefits_test.go

package invite

import (
	"context"
	"fmt"
	"strings"
	"testing"

	"github.com/DATA-DOG/go-sqlmock"
	"github.com/perfect-panel/server/pkg/tool"
	"gorm.io/driver/mysql"
	"gorm.io/gorm"
)

func TestQueryBenefitsCountsFamilyOwnerGiftAsInviteeGift(t *testing.T) {
	db, mock, cleanup := newBenefitsTestDB(t)
	defer cleanup()

	mock.ExpectQuery("COUNT(*) as cnt").
		WithArgs(int64(200), 2, 5).
		WillReturnRows(sqlmock.NewRows([]string{"user_id", "cnt"}).AddRow(200, 1))
	mock.ExpectQuery("SELECT user_id, subscription_user_id, order_no FROM `order`").
		WithArgs(int64(200), 2, 5).
		WillReturnRows(sqlmock.NewRows([]string{"user_id", "subscription_user_id", "order_no"}).AddRow(200, 900, "family-order"))
	mock.ExpectQuery("object_id IN").
		WithArgs(33, int64(100), "family-order", 331, 332).
		WillReturnRows(sqlmock.NewRows([]string{"object_id", "content"}))
	mock.ExpectQuery("object_id IN").
		WithArgs(34, int64(100), int64(200), int64(900), "family-order").
		WillReturnRows(sqlmock.NewRows([]string{"object_id", "content"}).
			AddRow(900, `{"type":341,"order_no":"family-order","amount":7,"balance":7,"remark":"邀请赠送"}`))

	benefits, err := QueryBenefits(context.Background(), db, []InviteRelation{{InviteeId: 200, InviterId: 100}})
	if err != nil {
		t.Fatalf("QueryBenefits returned error: %v", err)
	}

	benefit := benefits[200]
	if benefit.InviteeGiftDays != 7 {
		t.Fatalf("InviteeGiftDays = %d, want 7", benefit.InviteeGiftDays)
	}
	if benefit.InviterGiftDays != 0 {
		t.Fatalf("InviterGiftDays = %d, want 0", benefit.InviterGiftDays)
	}
	assertBenefitsExpectations(t, mock)
}

func TestQueryBenefitsKeepsDirectInviteeGift(t *testing.T) {
	db, mock, cleanup := newBenefitsTestDB(t)
	defer cleanup()

	mock.ExpectQuery("COUNT(*) as cnt").
		WithArgs(int64(200), 2, 5).
		WillReturnRows(sqlmock.NewRows([]string{"user_id", "cnt"}).AddRow(200, 1))
	mock.ExpectQuery("SELECT user_id, subscription_user_id, order_no FROM `order`").
		WithArgs(int64(200), 2, 5).
		WillReturnRows(sqlmock.NewRows([]string{"user_id", "subscription_user_id", "order_no"}).AddRow(200, 0, "direct-order"))
	mock.ExpectQuery("object_id IN").
		WithArgs(33, int64(100), "direct-order", 331, 332).
		WillReturnRows(sqlmock.NewRows([]string{"object_id", "content"}))
	mock.ExpectQuery("object_id IN").
		WithArgs(34, int64(100), int64(200), "direct-order").
		WillReturnRows(sqlmock.NewRows([]string{"object_id", "content"}).
			AddRow(200, `{"type":341,"order_no":"direct-order","amount":5,"balance":5,"remark":"邀请赠送"}`))

	benefits, err := QueryBenefits(context.Background(), db, []InviteRelation{{InviteeId: 200, InviterId: 100}})
	if err != nil {
		t.Fatalf("QueryBenefits returned error: %v", err)
	}

	if got := benefits[200].InviteeGiftDays; got != 5 {
		t.Fatalf("InviteeGiftDays = %d, want 5", got)
	}
	assertBenefitsExpectations(t, mock)
}

func TestQueryDeviceIdentifiersReturnsDeviceIdentifierAndNo(t *testing.T) {
	db, mock, cleanup := newBenefitsTestDB(t)
	defer cleanup()

	mock.ExpectQuery("FROM user_device ud").
		WithArgs(int64(100), int64(200)).
		WillReturnRows(sqlmock.NewRows([]string{"user_id", "device_id", "identifier"}).
			AddRow(100, 11, "inviter-hash").
			AddRow(200, 22, "invitee-hash"))

	identifiers, err := QueryDeviceIdentifiers(context.Background(), db, []int64{100, 200})
	if err != nil {
		t.Fatalf("QueryDeviceIdentifiers returned error: %v", err)
	}

	inviter := identifiers[100]
	if inviter.Identifier != "inviter-hash" {
		t.Fatalf("inviter identifier = %q, want inviter-hash", inviter.Identifier)
	}
	if inviter.DeviceNo != tool.DeviceIdToHash(11) {
		t.Fatalf("inviter device no = %q, want %q", inviter.DeviceNo, tool.DeviceIdToHash(11))
	}

	invitee := identifiers[200]
	if invitee.Identifier != "invitee-hash" {
		t.Fatalf("invitee identifier = %q, want invitee-hash", invitee.Identifier)
	}
	if invitee.DeviceNo != tool.DeviceIdToHash(22) {
		t.Fatalf("invitee device no = %q, want %q", invitee.DeviceNo, tool.DeviceIdToHash(22))
	}
	assertBenefitsExpectations(t, mock)
}

func newBenefitsTestDB(t *testing.T) (*gorm.DB, sqlmock.Sqlmock, func()) {
	t.Helper()

	sqlDB, mock, err := sqlmock.New(sqlmock.QueryMatcherOption(sqlmock.QueryMatcherFunc(func(expectedSQL, actualSQL string) error {
		if strings.Contains(actualSQL, expectedSQL) {
			return nil
		}
		return fmt.Errorf("actual sql %q does not contain %q", actualSQL, expectedSQL)
	})))
	if err != nil {
		t.Fatalf("create sqlmock: %v", err)
	}

	db, err := gorm.Open(mysql.New(mysql.Config{Conn: sqlDB, SkipInitializeWithVersion: true}), &gorm.Config{})
	if err != nil {
		_ = sqlDB.Close()
		t.Fatalf("open gorm db: %v", err)
	}

	return db, mock, func() {
		_ = sqlDB.Close()
	}
}

func assertBenefitsExpectations(t *testing.T, mock sqlmock.Sqlmock) {
	t.Helper()
	if err := mock.ExpectationsWereMet(); err != nil {
		t.Fatalf("unmet sql expectations: %v", err)
	}
}