new-api/middleware/distributor_concurrency_test.go

package middleware

import (
	"context"
	"errors"
	"net/http"
	"testing"
	"time"

	"github.com/QuantumNous/new-api/model"
	"github.com/QuantumNous/new-api/service"
	"github.com/QuantumNous/new-api/setting/operation_setting"
	"github.com/gin-gonic/gin"
	"github.com/stretchr/testify/require"
)

func restoreDistributorConcurrencyStubs(t *testing.T) {
	t.Helper()

	previousCheckQuota := checkChannelQuotaForDistribute
	previousAcquire := acquireChannelConcurrencyForDistribute
	previousSelect := cacheGetRandomSatisfiedChannelForDistribute
	previousWaitPlan := getChannelConcurrencyWaitPlanForDistribute
	previousEnterWaitQueue := tryEnterChannelConcurrencyWaitQueueForDistribute
	previousWaitForSlot := waitForChannelConcurrencySlotForDistribute
	t.Cleanup(func() {
		checkChannelQuotaForDistribute = previousCheckQuota
		acquireChannelConcurrencyForDistribute = previousAcquire
		cacheGetRandomSatisfiedChannelForDistribute = previousSelect
		getChannelConcurrencyWaitPlanForDistribute = previousWaitPlan
		tryEnterChannelConcurrencyWaitQueueForDistribute = previousEnterWaitQueue
		waitForChannelConcurrencySlotForDistribute = previousWaitForSlot
	})
}

func newDistributorTestContext() *gin.Context {
	gin.SetMode(gin.TestMode)
	c, _ := gin.CreateTestContext(nil)
	return c
}

func TestAcquireDistributedChannelConcurrencyReturns503WithoutWaitPlan(t *testing.T) {
	restoreDistributorConcurrencyStubs(t)

	checkChannelQuotaForDistribute = func(channelId int) error { return nil }
	acquireChannelConcurrencyForDistribute = func(channelId, limit int) (bool, func()) {
		return false, func() {}
	}
	cacheGetRandomSatisfiedChannelForDistribute = func(param *service.RetryParam) (*model.Channel, string, error) {
		return nil, param.TokenGroup, nil
	}
	getChannelConcurrencyWaitPlanForDistribute = func() operation_setting.ChannelConcurrencyWaitPlanSetting {
		return operation_setting.ChannelConcurrencyWaitPlanSetting{Enabled: false}
	}

	selected, release, response := acquireDistributedChannelConcurrency(
		newDistributorTestContext(),
		&model.Channel{Id: 1, ConcurrencyLimit: 1},
		"gpt-test",
		"default",
	)

	require.Nil(t, selected)
	require.Nil(t, release)
	require.NotNil(t, response)
	require.Equal(t, http.StatusServiceUnavailable, response.statusCode)
	require.Equal(t, "No available accounts", response.message)
}

func TestAcquireDistributedChannelConcurrencyReturns429WhenWaitQueueFull(t *testing.T) {
	restoreDistributorConcurrencyStubs(t)

	checkChannelQuotaForDistribute = func(channelId int) error { return nil }
	acquireChannelConcurrencyForDistribute = func(channelId, limit int) (bool, func()) {
		return false, func() {}
	}
	cacheGetRandomSatisfiedChannelForDistribute = func(param *service.RetryParam) (*model.Channel, string, error) {
		return nil, param.TokenGroup, nil
	}
	getChannelConcurrencyWaitPlanForDistribute = func() operation_setting.ChannelConcurrencyWaitPlanSetting {
		return operation_setting.ChannelConcurrencyWaitPlanSetting{
			Enabled:            true,
			TimeoutSeconds:     30,
			MaxWaitingRequests: 1,
		}
	}
	tryEnterChannelConcurrencyWaitQueueForDistribute = func(channelId, maxWaiting int) (bool, func()) {
		return false, func() {}
	}

	selected, release, response := acquireDistributedChannelConcurrency(
		newDistributorTestContext(),
		&model.Channel{Id: 1, ConcurrencyLimit: 1},
		"gpt-test",
		"default",
	)

	require.Nil(t, selected)
	require.Nil(t, release)
	require.NotNil(t, response)
	require.Equal(t, http.StatusTooManyRequests, response.statusCode)
	require.Equal(t, "Too many pending requests, please retry later", response.message)
}

func TestAcquireDistributedChannelConcurrencyWaitsForSlot(t *testing.T) {
	restoreDistributorConcurrencyStubs(t)

	released := false
	checkChannelQuotaForDistribute = func(channelId int) error { return nil }
	acquireChannelConcurrencyForDistribute = func(channelId, limit int) (bool, func()) {
		return false, func() {}
	}
	cacheGetRandomSatisfiedChannelForDistribute = func(param *service.RetryParam) (*model.Channel, string, error) {
		return nil, param.TokenGroup, nil
	}
	getChannelConcurrencyWaitPlanForDistribute = func() operation_setting.ChannelConcurrencyWaitPlanSetting {
		return operation_setting.ChannelConcurrencyWaitPlanSetting{
			Enabled:            true,
			TimeoutSeconds:     30,
			MaxWaitingRequests: 1,
		}
	}
	tryEnterChannelConcurrencyWaitQueueForDistribute = func(channelId, maxWaiting int) (bool, func()) {
		return true, func() {}
	}
	waitForChannelConcurrencySlotForDistribute = func(ctx context.Context, channelId, limit int, timeout time.Duration) (bool, func(), error) {
		return true, func() { released = true }, nil
	}

	selected, release, response := acquireDistributedChannelConcurrency(
		newDistributorTestContext(),
		&model.Channel{Id: 1, ConcurrencyLimit: 1},
		"gpt-test",
		"default",
	)

	require.Nil(t, response)
	require.NotNil(t, selected)
	require.Equal(t, 1, selected.Id)
	require.NotNil(t, release)
	release()
	require.True(t, released)
}

func TestAcquireDistributedChannelConcurrencyUsesAnotherAvailableChannel(t *testing.T) {
	restoreDistributorConcurrencyStubs(t)

	checkChannelQuotaForDistribute = func(channelId int) error {
		if channelId == 2 {
			return nil
		}
		return nil
	}
	acquireChannelConcurrencyForDistribute = func(channelId, limit int) (bool, func()) {
		if channelId == 2 {
			return true, func() {}
		}
		return false, func() {}
	}
	cacheGetRandomSatisfiedChannelForDistribute = func(param *service.RetryParam) (*model.Channel, string, error) {
		if _, excluded := param.ExcludedChannelIDs[1]; excluded {
			return &model.Channel{Id: 2, ConcurrencyLimit: 1}, param.TokenGroup, nil
		}
		return nil, param.TokenGroup, errors.New("expected saturated channel to be excluded")
	}
	getChannelConcurrencyWaitPlanForDistribute = func() operation_setting.ChannelConcurrencyWaitPlanSetting {
		return operation_setting.ChannelConcurrencyWaitPlanSetting{Enabled: false}
	}

	selected, release, response := acquireDistributedChannelConcurrency(
		newDistributorTestContext(),
		&model.Channel{Id: 1, ConcurrencyLimit: 1},
		"gpt-test",
		"default",
	)

	require.Nil(t, response)
	require.NotNil(t, selected)
	require.Equal(t, 2, selected.Id)
	require.NotNil(t, release)
	release()
}