new-api/middleware/concurrency_test.go

package middleware

import (
	"errors"
	"net/http"
	"net/http/httptest"
	"testing"

	"github.com/QuantumNous/new-api/model"

	"github.com/gin-gonic/gin"
	"github.com/stretchr/testify/require"
)

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

	previousGetUser := getUserByIdForConcurrency
	previousCheckRPM := checkUserRPMForConcurrency
	previousAcquire := acquireUserConcurrencyForCheck
	t.Cleanup(func() {
		getUserByIdForConcurrency = previousGetUser
		checkUserRPMForConcurrency = previousCheckRPM
		acquireUserConcurrencyForCheck = previousAcquire
	})
}

func performConcurrencyRequest(handler gin.HandlerFunc, userId int) *httptest.ResponseRecorder {
	gin.SetMode(gin.TestMode)
	router := gin.New()
	router.GET("/test", func(c *gin.Context) {
		if userId > 0 {
			c.Set("id", userId)
		}
	}, handler, func(c *gin.Context) {
		c.Status(http.StatusNoContent)
	})

	recorder := httptest.NewRecorder()
	request := httptest.NewRequest(http.MethodGet, "/test", nil)
	router.ServeHTTP(recorder, request)
	return recorder
}

func TestConcurrencyCheckAllowsWhenUserMissingFromContext(t *testing.T) {
	restoreConcurrencyStubs(t)

	getUserByIdForConcurrency = func(id int, selectAll bool) (*model.User, error) {
		t.Fatalf("expected user lookup to be skipped, got id %d", id)
		return nil, nil
	}

	recorder := performConcurrencyRequest(ConcurrencyCheck(), 0)

	require.Equal(t, http.StatusNoContent, recorder.Code)
}

func TestConcurrencyCheckFailOpenWhenUserLookupFails(t *testing.T) {
	restoreConcurrencyStubs(t)

	getUserByIdForConcurrency = func(id int, selectAll bool) (*model.User, error) {
		return nil, errors.New("database unavailable")
	}

	recorder := performConcurrencyRequest(ConcurrencyCheck(), 1)

	require.Equal(t, http.StatusNoContent, recorder.Code)
}

func TestConcurrencyCheckAllowsWhenLimitsDisabled(t *testing.T) {
	restoreConcurrencyStubs(t)

	getUserByIdForConcurrency = func(id int, selectAll bool) (*model.User, error) {
		return &model.User{Id: id}, nil
	}
	checkUserRPMForConcurrency = func(userId, limit int) bool {
		t.Fatalf("expected RPM check to be skipped")
		return false
	}
	acquireUserConcurrencyForCheck = func(userId, limit int) (bool, func()) {
		t.Fatalf("expected concurrency acquisition to be skipped")
		return false, func() {}
	}

	recorder := performConcurrencyRequest(ConcurrencyCheck(), 1)

	require.Equal(t, http.StatusNoContent, recorder.Code)
}

func TestConcurrencyCheckRejectsWhenRPMExceeded(t *testing.T) {
	restoreConcurrencyStubs(t)

	getUserByIdForConcurrency = func(id int, selectAll bool) (*model.User, error) {
		return &model.User{Id: id, RpmLimit: 1}, nil
	}
	checkUserRPMForConcurrency = func(userId, limit int) bool {
		return false
	}

	recorder := performConcurrencyRequest(ConcurrencyCheck(), 1)

	require.Equal(t, http.StatusTooManyRequests, recorder.Code)
}

func TestConcurrencyCheckRejectsWhenUserConcurrencyExceeded(t *testing.T) {
	restoreConcurrencyStubs(t)

	getUserByIdForConcurrency = func(id int, selectAll bool) (*model.User, error) {
		return &model.User{Id: id, ConcurrencyLimit: 1}, nil
	}
	acquireUserConcurrencyForCheck = func(userId, limit int) (bool, func()) {
		return false, func() {}
	}

	recorder := performConcurrencyRequest(ConcurrencyCheck(), 1)

	require.Equal(t, http.StatusTooManyRequests, recorder.Code)
}

func TestConcurrencyCheckReleasesAfterRequest(t *testing.T) {
	restoreConcurrencyStubs(t)

	released := false
	getUserByIdForConcurrency = func(id int, selectAll bool) (*model.User, error) {
		return &model.User{Id: id, ConcurrencyLimit: 1}, nil
	}
	acquireUserConcurrencyForCheck = func(userId, limit int) (bool, func()) {
		return true, func() {
			released = true
		}
	}

	recorder := performConcurrencyRequest(ConcurrencyCheck(), 1)

	require.Equal(t, http.StatusNoContent, recorder.Code)
	require.True(t, released)
}