feat: channel test (#6917)

* feat: channel test

* fix: code smell

controller/channel_test_internal_test.go

@@ -2,9 +2,11 @@ package controller
 
 import (
 	"bytes"
+	"context"
 	"fmt"
 	"net/http"
 	"net/http/httptest"
+	"sync/atomic"
 	"testing"
 
 	"github.com/QuantumNous/new-api/common"
@@ -339,6 +341,111 @@ func TestSelectChannelsForAutomaticTestAutoBanOnlyUsesEligibleChannels(t *testin
 	require.Equal(t, 3, selected[1].Id)
 }
 
+func TestRunChannelTestWorkersHonorsConfiguredConcurrency(t *testing.T) {
+	originalInterval := common.RequestInterval
+	common.RequestInterval = 0
+	t.Cleanup(func() { common.RequestInterval = originalInterval })
+
+	channels := []*model.Channel{
+		{Id: 1, Status: common.ChannelStatusEnabled},
+		{Id: 2, Status: common.ChannelStatusEnabled},
+		{Id: 3, Status: common.ChannelStatusEnabled},
+		{Id: 4, Status: common.ChannelStatusEnabled},
+	}
+	started := make(chan struct{}, len(channels))
+	release := make(chan struct{})
+	var active atomic.Int32
+	var maxActive atomic.Int32
+	progress := make([]int, 0, len(channels)+1)
+	summaryResult := make(chan channelTestSummary, 1)
+
+	go func() {
+		summaryResult <- runChannelTestWorkers(
+			context.Background(),
+			channels,
+			2,
+			func(_ context.Context, _ *model.Channel) channelTestSummary {
+				current := active.Add(1)
+				defer active.Add(-1)
+				for {
+					observed := maxActive.Load()
+					if current <= observed || maxActive.CompareAndSwap(observed, current) {
+						break
+					}
+				}
+				started <- struct{}{}
+				<-release
+				return channelTestSummary{Tested: 1, Succeeded: 1}
+			},
+			func(processed, _ int) {
+				progress = append(progress, processed)
+			},
+		)
+	}()
+
+	<-started
+	<-started
+	select {
+	case <-started:
+		t.Fatal("started more channel tests than the configured concurrency")
+	default:
+	}
+	close(release)
+
+	summary := <-summaryResult
+
+	assert.Equal(t, int32(2), maxActive.Load())
+	assert.Equal(t, channelTestSummary{Tested: 4, Succeeded: 4}, summary)
+	assert.Equal(t, []int{0, 1, 2, 3, 4}, progress)
+}
+
+func TestRunChannelTestWorkersStopsAfterCancellation(t *testing.T) {
+	originalInterval := common.RequestInterval
+	common.RequestInterval = 0
+	t.Cleanup(func() { common.RequestInterval = originalInterval })
+
+	ctx, cancel := context.WithCancel(context.Background())
+	channels := []*model.Channel{
+		{Id: 1, Status: common.ChannelStatusEnabled},
+		{Id: 2, Status: common.ChannelStatusEnabled},
+		{Id: 3, Status: common.ChannelStatusEnabled},
+		{Id: 4, Status: common.ChannelStatusEnabled},
+	}
+	started := make(chan struct{}, len(channels))
+	progress := make([]int, 0, 1)
+	summaryResult := make(chan channelTestSummary, 1)
+
+	go func() {
+		summaryResult <- runChannelTestWorkers(
+			ctx,
+			channels,
+			2,
+			func(ctx context.Context, _ *model.Channel) channelTestSummary {
+				started <- struct{}{}
+				<-ctx.Done()
+				return channelTestSummary{Tested: 1, Succeeded: 1}
+			},
+			func(processed, _ int) {
+				progress = append(progress, processed)
+			},
+		)
+	}()
+
+	<-started
+	<-started
+	cancel()
+
+	summary := <-summaryResult
+
+	select {
+	case <-started:
+		t.Fatal("started another channel test after cancellation")
+	default:
+	}
+	assert.Equal(t, channelTestSummary{Tested: 2, Succeeded: 2}, summary)
+	assert.Equal(t, []int{0}, progress)
+}
+
 func TestTestAllChannelsRejectsExistingActiveTask(t *testing.T) {
 	db := setupModelListControllerTestDB(t)
 	require.NoError(t, db.AutoMigrate(&model.SystemTask{}, &model.SystemTaskLock{}))