package invite import ( "context" "fmt" "strings" "testing" "github.com/DATA-DOG/go-sqlmock" "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(900), int64(200), int64(100), "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(200), int64(100), "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 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) } }