3x-ui/internal/web/service/inbound_client_stats_test.go

package service

import (
	"errors"
	"fmt"
	"reflect"
	"testing"

	"gorm.io/gorm"

	"github.com/mhsanaei/3x-ui/v3/internal/database"
	"github.com/mhsanaei/3x-ui/v3/internal/database/model"
	"github.com/mhsanaei/3x-ui/v3/internal/xray"
)

func seedStatsIdentity(t testing.TB, inboundIDs []int, email string) model.ClientRecord {
	t.Helper()
	db := database.GetDB()
	rec := model.ClientRecord{Email: email, UUID: "uuid-" + email, SubID: "sub-" + email, Enable: true}
	if err := db.Create(&rec).Error; err != nil {
		t.Fatal(err)
	}
	for _, id := range inboundIDs {
		if err := db.Create(&model.ClientInbound{ClientId: rec.Id, InboundId: id}).Error; err != nil {
			t.Fatal(err)
		}
	}
	return rec
}

func TestClientStatsIdentitiesBatchesPreloadedInbounds(t *testing.T) {
	setupBulkDB(t)
	var inbounds []*model.Inbound
	for i := range 8 {
		ib := mkInbound(t, 57000+i, model.VLESS, `{"clients":[]}`)
		rec := seedStatsIdentity(t, []int{ib.Id}, fmt.Sprintf("batch-%d", i))
		ib.ClientStats = []xray.ClientTraffic{{Email: rec.Email, Up: 123, Down: 456, UUID: "stale", SubId: "stale"}}
		inbounds = append(inbounds, ib)
	}
	queries := countClientTableQueries(t, "stats_batch", func() { (&InboundService{}).enrichClientStats(database.GetDB(), inbounds) })
	if queries != 1 {
		t.Fatalf("clients queries = %d, want 1 for eight preloaded inbounds", queries)
	}
	for _, ib := range inbounds {
		st := ib.ClientStats[0]
		if st.UUID != "uuid-"+st.Email || st.SubId != "sub-"+st.Email || st.Up != 123 || st.Down != 456 {
			t.Fatalf("enriched stats = %+v", st)
		}
	}
}

func TestBackfillClientStatsChunkBoundaryAndEmpty(t *testing.T) {
	setupBulkDB(t)
	var inbounds []*model.Inbound
	for i := range 401 {
		inbounds = append(inbounds, &model.Inbound{Id: i + 1})
	}
	seedStatsIdentity(t, []int{1, 401}, "boundary")
	if err := database.GetDB().Create(&xray.ClientTraffic{Email: "boundary", InboundId: 999, Up: 7}).Error; err != nil {
		t.Fatal(err)
	}
	inbounds = append(inbounds, &model.Inbound{Id: 1}, &model.Inbound{Id: 401})
	queries := countClientTableQueries(t, "stats_boundary", func() { (&InboundService{}).enrichClientStats(database.GetDB(), inbounds) })
	if queries != 2 {
		t.Fatalf("clients queries = %d, want 2 for 401 distinct IDs plus duplicates", queries)
	}
	for _, i := range []int{0, 400} {
		if len(inbounds[i].ClientStats) != 1 || inbounds[i].ClientStats[0].Up != 7 {
			t.Fatalf("inbound %d stats = %+v", i, inbounds[i].ClientStats)
		}
	}
	queries = countClientTableQueries(t, "stats_empty", func() { (&InboundService{}).backfillClientStats(database.GetDB(), nil) })
	if queries != 0 {
		t.Fatalf("empty request queries = %d", queries)
	}
}

func TestBackfillClientStatsProjectionFailurePreservesStats(t *testing.T) {
	setupBulkDB(t)
	ib := &model.Inbound{Id: 1, ClientStats: []xray.ClientTraffic{{Email: "preloaded", UUID: "old", SubId: "old-sub", Up: 5}}}
	seedStatsIdentity(t, []int{1}, "preloaded")
	seedStatsIdentity(t, []int{1}, "missing")
	db := database.GetDB()
	if err := db.Create(&xray.ClientTraffic{Email: "missing", InboundId: 2, Up: 9}).Error; err != nil {
		t.Fatal(err)
	}
	want := append([]xray.ClientTraffic(nil), ib.ClientStats...)
	injected := errors.New("stats projection unavailable")
	name := "test:stats_projection_failure"
	queries := 0
	if err := db.Callback().Query().Before("gorm:query").Register(name, func(tx *gorm.DB) {
		if tx.Statement.Table == "clients" {
			queries++
			if queries == 2 {
				tx.AddError(injected)
			}
		}
	}); err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = db.Callback().Query().Remove(name) })
	inbounds := []*model.Inbound{ib}
	for id := 2; id <= 401; id++ {
		inbounds = append(inbounds, &model.Inbound{Id: id})
	}
	(&InboundService{}).enrichClientStats(db, inbounds)
	if !reflect.DeepEqual(ib.ClientStats, want) {
		t.Fatalf("projection failure mutated stats: got %+v, want %+v", ib.ClientStats, want)
	}
}

func TestBackfillClientStatsTrafficPageFailureKeepsFreshPreload(t *testing.T) {
	setupBulkDB(t)
	db := database.GetDB()
	ib := &model.Inbound{Id: 1, ClientStats: []xray.ClientTraffic{{Email: "preloaded", UUID: "old", SubId: "old-sub", Up: 5}}}
	seedStatsIdentity(t, []int{1}, "preloaded")
	for i := range 401 {
		email := fmt.Sprintf("page-%03d", i)
		seedStatsIdentity(t, []int{1}, email)
		if err := db.Create(&xray.ClientTraffic{Email: email, InboundId: 2, Up: int64(i)}).Error; err != nil {
			t.Fatal(err)
		}
	}
	pages := 0
	name := "test:stats_traffic_page_failure"
	if err := db.Callback().Query().Before("gorm:query").Register(name, func(tx *gorm.DB) {
		if tx.Statement.Table == "client_traffics" {
			pages++
			if pages == 2 {
				tx.AddError(errors.New("second traffic page unavailable"))
			}
		}
	}); err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = db.Callback().Query().Remove(name) })
	(&InboundService{}).enrichClientStats(db, []*model.Inbound{ib})
	want := []xray.ClientTraffic{{Email: "preloaded", UUID: "uuid-preloaded", SubId: "sub-preloaded", Up: 5}}
	if pages != 2 || !reflect.DeepEqual(ib.ClientStats, want) {
		t.Fatalf("pages=%d stats=%+v, want two pages and %+v without partial append", pages, ib.ClientStats, want)
	}
}

func TestEnrichClientStatsPreservesSiblingTrafficAndClientOrder(t *testing.T) {
	setupBulkDB(t)
	db := database.GetDB()
	for _, email := range []string{"second", "First", "absent", "FIRST"} {
		ids := []int{12}
		if email == "second" {
			ids = append(ids, 13)
		}
		seedStatsIdentity(t, ids, email)
	}
	rows := []xray.ClientTraffic{
		{Email: "First", InboundId: 999, Enable: true, Up: 123, Down: 456, Total: 789, ExpiryTime: 12345, LastOnline: 2468, Reset: 3},
		{Email: "second", InboundId: 999, Enable: false, Up: 10, Down: 20},
	}
	if err := db.Create(&rows).Error; err != nil {
		t.Fatal(err)
	}
	prefix := xray.ClientTraffic{Email: "unlinked-prefix", Up: 77}
	inbounds := []*model.Inbound{{Id: 12, ClientStats: []xray.ClientTraffic{prefix}}, {Id: 13}, {Id: 14}}
	for range 2 {
		(&InboundService{}).enrichClientStats(db, inbounds)
	}
	wantFirst := rows[0]
	wantFirst.UUID = "uuid-FIRST"
	wantFirst.SubId = "sub-FIRST"
	wantSecond := rows[1]
	wantSecond.UUID = "uuid-second"
	wantSecond.SubId = "sub-second"
	if want := []xray.ClientTraffic{prefix, wantSecond, wantFirst}; !reflect.DeepEqual(inbounds[0].ClientStats, want) {
		t.Fatalf("stats = %+v, want %+v", inbounds[0].ClientStats, want)
	}
	if !reflect.DeepEqual(inbounds[1].ClientStats, []xray.ClientTraffic{wantSecond}) || len(inbounds[2].ClientStats) != 0 {
		t.Fatalf("shared/empty stats = %+v / %+v", inbounds[1].ClientStats, inbounds[2].ClientStats)
	}
	if err := db.Model(&model.ClientRecord{}).Where("email = ?", "FIRST").Updates(map[string]any{"uuid": "", "sub_id": ""}).Error; err != nil {
		t.Fatal(err)
	}
	(&InboundService{}).enrichClientStats(db, inbounds)
	if st := inbounds[0].ClientStats[2]; st.UUID != "" || st.SubId != "" {
		t.Fatalf("empty identity failed to clear stale values: %+v", st)
	}
}

func TestClientStatsIdentitiesUsesSuppliedTransaction(t *testing.T) {
	setupBulkDB(t)
	db := database.GetDB()
	rec := seedStatsIdentity(t, []int{1}, "transaction")
	tx := db.Begin()
	if tx.Error != nil {
		t.Fatal(tx.Error)
	}
	defer tx.Rollback()
	if err := tx.Model(&rec).Updates(map[string]any{"uuid": "uncommitted", "sub_id": "uncommitted-sub"}).Error; err != nil {
		t.Fatal(err)
	}
	ib := &model.Inbound{Id: 1, ClientStats: []xray.ClientTraffic{{Email: rec.Email}}}
	(&InboundService{}).enrichClientStats(tx, []*model.Inbound{ib})
	if st := ib.ClientStats[0]; st.UUID != "uncommitted" || st.SubId != "uncommitted-sub" {
		t.Fatalf("transaction enrichment = %+v, want uncommitted identity", st)
	}
	if err := tx.Rollback().Error; err != nil {
		t.Fatal(err)
	}
	identities, err := (&ClientService{}).listStatsIdentitiesForInbounds(nil, []int{1, 1})
	if err != nil {
		t.Fatal(err)
	}
	want := []clientStatsIdentity{{UUID: "uuid-transaction", Email: "transaction", SubID: "sub-transaction"}}
	if !reflect.DeepEqual(identities[1], want) {
		t.Fatalf("fallback after rollback = %+v, want %+v", identities[1], want)
	}
}