diff --git a/core/http/admission_handler_test.go b/core/http/admission_handler_test.go index 77ec65da0..6a0a2201d 100644 --- a/core/http/admission_handler_test.go +++ b/core/http/admission_handler_test.go @@ -5,16 +5,17 @@ import ( "fmt" "net/http" "net/http/httptest" - "testing" "time" "github.com/labstack/echo/v4" corebackend "github.com/mudler/LocalAI/core/backend" "github.com/mudler/LocalAI/core/services/nodes" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" ) -func TestApplyBackendAdmission(t *testing.T) { - t.Run("maps BackendAdmissionError to 429 with Retry-After", func(t *testing.T) { +var _ = Describe("Backend admission", func() { + It("maps BackendAdmissionError to 429 with Retry-After", func() { e := echo.New() req := httptest.NewRequest(http.MethodPost, "/", nil) rec := httptest.NewRecorder() @@ -23,61 +24,47 @@ func TestApplyBackendAdmission(t *testing.T) { err := &corebackend.BackendAdmissionError{Limit: 4, RetryAfter: 3 * time.Second} code := applyBackendAdmission(err, http.StatusInternalServerError, c) - if code != http.StatusTooManyRequests { - t.Fatalf("expected 429, got %d", code) - } - if got := rec.Header().Get("Retry-After"); got != "3" { - t.Fatalf("expected Retry-After 3, got %q", got) - } + Expect(code).To(Equal(http.StatusTooManyRequests)) + Expect(rec.Header().Get("Retry-After")).To(Equal("3")) }) - t.Run("passes through non-admission errors unchanged", func(t *testing.T) { + It("passes through non-admission errors unchanged", func() { e := echo.New() req := httptest.NewRequest(http.MethodPost, "/", nil) rec := httptest.NewRecorder() c := e.NewContext(req, rec) code := applyBackendAdmission(errors.New("some other error"), http.StatusInternalServerError, c) - if code != http.StatusInternalServerError { - t.Fatalf("expected 500, got %d", code) - } - if got := rec.Header().Get("Retry-After"); got != "" { - t.Fatalf("expected no Retry-After, got %q", got) - } + Expect(code).To(Equal(http.StatusInternalServerError)) + Expect(rec.Header().Get("Retry-After")).To(BeEmpty()) }) -} +}) -func TestApplyNoAvailableNodes(t *testing.T) { - t.Run("maps ErrNoAvailableNodes to 503", func(t *testing.T) { +var _ = Describe("No available nodes", func() { + It("maps ErrNoAvailableNodes to 503", func() { // The scheduler wraps the sentinel in fmt.Errorf chains and via // errors.Join — errors.Is must still find it. wrapped := fmt.Errorf("routing model foo: %w", fmt.Errorf("no available nodes: %w", fmt.Errorf("no healthy nodes available: %w", - errors.Join(nodes.ErrEvictionBusy, nodes.ErrNoAvailableNodes)))) + errors.Join(nodes.ErrEvictionBusy, nodes.ErrNoAvailableNodes)))) code := applyNoAvailableNodes(wrapped, http.StatusInternalServerError) - if code != http.StatusServiceUnavailable { - t.Fatalf("expected 503, got %d", code) - } + Expect(code).To(Equal(http.StatusServiceUnavailable)) }) - t.Run("maps selector-mismatch chain to 503", func(t *testing.T) { + It("maps selector-mismatch chain to 503", func() { wrapped := fmt.Errorf("routing model bar: %w", fmt.Errorf("no available nodes: %w", fmt.Errorf("no healthy nodes match selector for model bar: {\"gpu.vendor\":\"tpu\"}: %w", nodes.ErrNoAvailableNodes))) code := applyNoAvailableNodes(wrapped, http.StatusInternalServerError) - if code != http.StatusServiceUnavailable { - t.Fatalf("expected 503, got %d", code) - } + Expect(code).To(Equal(http.StatusServiceUnavailable)) }) - t.Run("passes through unrelated errors unchanged", func(t *testing.T) { + It("passes through unrelated errors unchanged", func() { code := applyNoAvailableNodes(errors.New("database timeout"), http.StatusInternalServerError) - if code != http.StatusInternalServerError { - t.Fatalf("expected 500, got %d", code) - } + Expect(code).To(Equal(http.StatusInternalServerError)) }) -} +})