mirror of
https://github.com/mudler/LocalAI.git
synced 2026-11-06 12:33:16 -05:00
Ettore Di Giacinto
b15905e90f
The previous commit was made on a stale base and so reverted work that landed on master before it. This commit puts every file back to the state of its parent's parent. The permission enforcement of the previous commit stays. Assisted-by: Claude:claude-sonnet-5-5 go
45 lines
1.7 KiB
Go
45 lines
1.7 KiB
Go
// SPDX-License-Identifier: MIT
|
|
package protocol_test
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"time"
|
|
|
|
"github.com/mudler/LocalAI/addons/internal/idle"
|
|
. "github.com/mudler/LocalAI/addons/internal/protocol"
|
|
"github.com/mudler/LocalAI/pkg/addons/contract"
|
|
. "github.com/onsi/ginkgo/v2"
|
|
. "github.com/onsi/gomega"
|
|
)
|
|
|
|
type limitedRuntime struct{ idle.Runtime }
|
|
|
|
func (limitedRuntime) Read(context.Context, contract.Envelope) (contract.Result, error) {
|
|
return contract.Result{}, fmt.Errorf("service: %w", &contract.ErrRateLimited{RetryAfter: 61})
|
|
}
|
|
|
|
var _ = Describe("Handler and a rate limited read", func() {
|
|
It("answers 429 with the wait in Retry-After and a fixed body", func() {
|
|
startup := Startup{Registration: Registration{AddonID: "sample", Generation: 1, Incarnation: "inc", Protocol: 1}, Credential: "secret", Limits: contract.DefaultLimits()}
|
|
env := contract.Envelope{
|
|
Request: contract.Request{AddonID: "sample", OperationID: "op", Generation: 1, Action: "read", Arguments: json.RawMessage(`{}`)},
|
|
Authority: contract.Authority{Actor: contract.Actor{UserID: "u", CredentialID: "c", Mode: "named-user"}, Scope: "read", ExpiresAt: time.Now().Add(time.Minute)},
|
|
}
|
|
b, _ := json.Marshal(env)
|
|
req := httptest.NewRequest(http.MethodPost, "/read", bytes.NewReader(b))
|
|
req.Header.Set("Authorization", "Bearer secret")
|
|
req.Header.Set("X-LocalAI-Incarnation", "inc")
|
|
req.Header.Set("Content-Type", "application/json")
|
|
rec := httptest.NewRecorder()
|
|
Handler(startup, limitedRuntime{}).ServeHTTP(rec, req)
|
|
Expect(rec.Code).To(Equal(http.StatusTooManyRequests))
|
|
Expect(rec.Header().Get("Retry-After")).To(Equal("61"))
|
|
Expect(rec.Body.String()).To(Equal("rate_limited\n"))
|
|
})
|
|
})
|