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

package service

import (
	"fmt"
	"sync/atomic"
	"testing"

	"gorm.io/gorm"

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

// countClientTableQueries runs fn with a callback counting SELECTs against the
// clients table, so a per-email lookup shows up as growth with the batch size.
func countClientTableQueries(t *testing.T, name string, fn func()) int {
	t.Helper()
	db := database.GetDB()
	var n int64
	cb := "test:count_clients_query_" + name
	if err := db.Callback().Query().Before("gorm:query").Register(cb, func(tx *gorm.DB) {
		if tx.Statement != nil && tx.Statement.Table == "clients" {
			atomic.AddInt64(&n, 1)
		}
	}); err != nil {
		t.Fatalf("register query callback: %v", err)
	}
	defer func() {
		if err := db.Callback().Query().Remove(cb); err != nil {
			t.Errorf("remove query callback: %v", err)
		}
	}()
	fn()
	return int(atomic.LoadInt64(&n))
}

func seedEnabledClientsForReset(t *testing.T, svc *ClientService, port int, n int, prefix string) []string {
	t.Helper()
	clients := make([]model.Client, 0, n)
	for i := range n {
		email := fmt.Sprintf("%s-%d@x", prefix, i)
		clients = append(clients, model.Client{
			Email:  email,
			ID:     fmt.Sprintf("%08d-1111-1111-1111-111111111111", i),
			SubID:  email,
			Enable: true,
		})
	}
	ib := mkInbound(t, port, model.VLESS, clientsSettings(t, clients))
	if err := svc.SyncInbound(nil, ib.Id, clients); err != nil {
		t.Fatalf("seed linkage: %v", err)
	}
	emails := make([]string, 0, n)
	for _, c := range clients {
		mkTraffic(t, ib.Id, c.Email, 100, 200, 0, 0, true)
		emails = append(emails, c.Email)
	}
	return emails
}

// TestBulkResetTraffic_DoesNotQueryPerEmail pins BulkResetTraffic's client
// lookup to a batched read: the number of SELECTs on clients must not grow
// with the number of emails reset.
func TestBulkResetTraffic_DoesNotQueryPerEmail(t *testing.T) {
	setupBulkDB(t)
	svc := &ClientService{}
	inboundSvc := &InboundService{}

	few := seedEnabledClientsForReset(t, svc, 53010, 3, "few")
	many := seedEnabledClientsForReset(t, svc, 53011, 30, "many")

	fewCount := countClientTableQueries(t, "few", func() {
		if _, err := svc.BulkResetTraffic(inboundSvc, few); err != nil {
			t.Fatalf("BulkResetTraffic(few): %v", err)
		}
	})
	manyCount := countClientTableQueries(t, "many", func() {
		if _, err := svc.BulkResetTraffic(inboundSvc, many); err != nil {
			t.Fatalf("BulkResetTraffic(many): %v", err)
		}
	})

	if manyCount != fewCount {
		t.Fatalf("clients SELECTs: %d emails -> %d, %d emails -> %d; want the same batched count",
			len(few), fewCount, len(many), manyCount)
	}
}