3a25007df2
Add user RPM and concurrency middleware, enforce channel concurrency for relay selection and retry paths, and wire checks into relay routes.
78 lines
1.8 KiB
Go
78 lines
1.8 KiB
Go
package controller
|
|
|
|
import (
|
|
"net/http"
|
|
"testing"
|
|
|
|
"github.com/QuantumNous/new-api/model"
|
|
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
func restoreRelayConcurrencyStub(t *testing.T) {
|
|
t.Helper()
|
|
|
|
previousAcquire := acquireChannelConcurrencyForRelay
|
|
t.Cleanup(func() {
|
|
acquireChannelConcurrencyForRelay = previousAcquire
|
|
})
|
|
}
|
|
|
|
func TestAcquireRetryChannelConcurrencySkipsInitialAttempt(t *testing.T) {
|
|
restoreRelayConcurrencyStub(t)
|
|
|
|
acquireChannelConcurrencyForRelay = func(channelId, limit int) (bool, func()) {
|
|
t.Fatalf("expected initial attempt to use distributor-owned acquisition")
|
|
return false, func() {}
|
|
}
|
|
|
|
release, err := acquireRetryChannelConcurrency(&model.Channel{
|
|
Id: 1,
|
|
ConcurrencyLimit: 1,
|
|
}, 0)
|
|
|
|
require.Nil(t, err)
|
|
require.NotNil(t, release)
|
|
release()
|
|
}
|
|
|
|
func TestAcquireRetryChannelConcurrencyRejectsExceededRetryChannel(t *testing.T) {
|
|
restoreRelayConcurrencyStub(t)
|
|
|
|
acquireChannelConcurrencyForRelay = func(channelId, limit int) (bool, func()) {
|
|
require.Equal(t, 2, channelId)
|
|
require.Equal(t, 1, limit)
|
|
return false, func() {}
|
|
}
|
|
|
|
release, err := acquireRetryChannelConcurrency(&model.Channel{
|
|
Id: 2,
|
|
ConcurrencyLimit: 1,
|
|
}, 1)
|
|
|
|
require.Nil(t, release)
|
|
require.NotNil(t, err)
|
|
require.Equal(t, http.StatusTooManyRequests, err.StatusCode)
|
|
}
|
|
|
|
func TestAcquireRetryChannelConcurrencyReleasesRetryChannel(t *testing.T) {
|
|
restoreRelayConcurrencyStub(t)
|
|
|
|
released := false
|
|
acquireChannelConcurrencyForRelay = func(channelId, limit int) (bool, func()) {
|
|
return true, func() {
|
|
released = true
|
|
}
|
|
}
|
|
|
|
release, err := acquireRetryChannelConcurrency(&model.Channel{
|
|
Id: 3,
|
|
ConcurrencyLimit: 1,
|
|
}, 1)
|
|
|
|
require.Nil(t, err)
|
|
require.NotNil(t, release)
|
|
release()
|
|
require.True(t, released)
|
|
}
|