nex/backend/tests/mysql/concurrent_test.go

//go:build mysql

package mysql

import (
	"sync"
	"testing"

	"github.com/google/uuid"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"

	"nex/backend/internal/config"
	"nex/backend/internal/repository"
)

func TestConcurrent_UsageStatsRecord(t *testing.T) {
	db := SetupMySQLTestDB(t)

	statsRepo := repository.NewStatsRepository(db)

	providerID := "concurrent-test-provider"
	modelName := "gpt-4"

	concurrency := 10
	var wg sync.WaitGroup
	wg.Add(concurrency)

	errChan := make(chan error, concurrency)

	for i := 0; i < concurrency; i++ {
		go func() {
			defer wg.Done()
			err := statsRepo.Record(providerID, modelName)
			if err != nil {
				errChan <- err
			}
		}()
	}

	wg.Wait()
	close(errChan)

	var errorCount int
	uniqueErrors := make(map[string]int)
	for err := range errChan {
		errorCount++
		uniqueErrors[err.Error()]++
	}

	t.Logf("并发 %d 次，错误 %d 次", concurrency, errorCount)
	for errMsg, count := range uniqueErrors {
		t.Logf("  错误: %s (出现 %d 次)", errMsg, count)
	}

	var stats config.UsageStats
	err := db.Where("provider_id = ? AND model_name = ?", providerID, modelName).
		First(&stats).Error
	require.NoError(t, err, "应能查到 usage_stats 记录")

	successCount := concurrency - errorCount
	t.Logf("成功次数: %d, 最终 request_count: %d", successCount, stats.RequestCount)

	assert.Equal(t, concurrency, stats.RequestCount, "request_count 应等于并发数，无数据丢失或重复")
}

func TestConcurrent_ProviderCreate(t *testing.T) {
	db := SetupMySQLTestDB(t)

	providerID := "concurrent-provider-id"
	concurrency := 10

	var wg sync.WaitGroup
	wg.Add(concurrency)

	successCount := 0
	var mu sync.Mutex

	for i := 0; i < concurrency; i++ {
		go func() {
			defer wg.Done()

			provider := config.Provider{
				ID:      providerID,
				Name:    "Concurrent Provider",
				APIKey:  "test-key",
				BaseURL: "https://test.com",
				Enabled: true,
			}

			err := db.Create(&provider).Error
			if err == nil {
				mu.Lock()
				successCount++
				mu.Unlock()
			}
		}()
	}

	wg.Wait()

	assert.Equal(t, 1, successCount, "仅 1 个创建应成功")

	var count int64
	db.Model(&config.Provider{}).Where("id = ?", providerID).Count(&count)
	assert.Equal(t, int64(1), count, "最终应有 1 条记录")
}

func TestConcurrent_ModelCreate(t *testing.T) {
	db := SetupMySQLTestDB(t)

	provider := config.Provider{
		ID:      "concurrent-model-provider",
		Name:    "Test Provider",
		APIKey:  "test-key",
		BaseURL: "https://test.com",
		Enabled: true,
	}
	err := db.Create(&provider).Error
	require.NoError(t, err, "创建 provider 应成功")

	modelName := "gpt-4-concurrent"
	concurrency := 10

	var wg sync.WaitGroup
	wg.Add(concurrency)

	successCount := 0
	var mu sync.Mutex

	for i := 0; i < concurrency; i++ {
		go func(idx int) {
			defer wg.Done()

			model := config.Model{
				ID:         uuid.New().String(),
				ProviderID: provider.ID,
				ModelName:  modelName,
				Enabled:    true,
			}

			err := db.Create(&model).Error
			if err == nil {
				mu.Lock()
				successCount++
				mu.Unlock()
			}
		}(i)
	}

	wg.Wait()

	assert.Equal(t, 1, successCount, "仅 1 个创建应成功")

	var count int64
	db.Model(&config.Model{}).Where("provider_id = ? AND model_name = ?", provider.ID, modelName).Count(&count)
	assert.Equal(t, int64(1), count, "最终应有 1 条记录")
}