71 lines
3.3 KiB
Go
71 lines
3.3 KiB
Go
package main
|
|
|
|
import (
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
|
|
"git.cin.su/m/Assistant/contracts/capability"
|
|
)
|
|
|
|
func TestFleetReadinessSelectsBestAvailablePrimaryTierBeforeLifeboat(t *testing.T) {
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { w.WriteHeader(http.StatusOK) }))
|
|
defer server.Close()
|
|
handler := fleetReadinessHandler(server.Client(), capability.FleetTierHeavy, []fleetEndpoint{
|
|
{ProviderID: "cpu", ModelID: "cpu-model", Tier: capability.FleetTierFast, URL: server.URL},
|
|
{ProviderID: "medium", ModelID: "medium-model", Tier: capability.FleetTierMedium},
|
|
{ProviderID: "heavy", ModelID: "heavy-model", Tier: capability.FleetTierHeavy},
|
|
{ProviderID: "lifeboat", ModelID: "lifeboat-model", Tier: capability.FleetTierLifeboat, URL: server.URL},
|
|
})
|
|
recorder := httptest.NewRecorder()
|
|
request := httptest.NewRequest(http.MethodGet, "/v1/fleet/readiness", nil)
|
|
handler.ServeHTTP(recorder, request)
|
|
if recorder.Code != http.StatusOK || !strings.Contains(recorder.Body.String(), `"selected_tier":"FAST"`) || !strings.Contains(recorder.Body.String(), `"tier":"MEDIUM"`) {
|
|
t.Fatalf("status=%d body=%s", recorder.Code, recorder.Body.String())
|
|
}
|
|
}
|
|
|
|
func TestFleetReadinessDoesNotWaitForSlowMembersSerially(t *testing.T) {
|
|
slow := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
|
time.Sleep(150 * time.Millisecond)
|
|
w.WriteHeader(http.StatusServiceUnavailable)
|
|
}))
|
|
defer slow.Close()
|
|
fast := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { w.WriteHeader(http.StatusOK) }))
|
|
defer fast.Close()
|
|
handler := fleetReadinessHandler(fast.Client(), capability.FleetTierFast, []fleetEndpoint{
|
|
{ProviderID: "slow", ModelID: "slow-model", Tier: capability.FleetTierMedium, URL: slow.URL},
|
|
{ProviderID: "lifeboat", ModelID: "lifeboat-model", Tier: capability.FleetTierLifeboat, URL: fast.URL},
|
|
})
|
|
started := time.Now()
|
|
recorder := httptest.NewRecorder()
|
|
handler.ServeHTTP(recorder, httptest.NewRequest(http.MethodGet, "/v1/fleet/readiness", nil))
|
|
if recorder.Code != http.StatusOK || !strings.Contains(recorder.Body.String(), `"selected_tier":"LIFEBOAT"`) {
|
|
t.Fatalf("status=%d body=%s", recorder.Code, recorder.Body.String())
|
|
}
|
|
if elapsed := time.Since(started); elapsed >= 300*time.Millisecond {
|
|
t.Fatalf("fleet readiness was serialized: elapsed=%s", elapsed)
|
|
}
|
|
}
|
|
|
|
func TestFleetReadinessSendsConfiguredProviderHeaders(t *testing.T) {
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
if r.Header.Get("Authorization") != "Bearer fleet-key" {
|
|
t.Fatalf("authorization=%q", r.Header.Get("Authorization"))
|
|
}
|
|
w.WriteHeader(http.StatusOK)
|
|
}))
|
|
defer server.Close()
|
|
handler := fleetReadinessHandler(server.Client(), capability.FleetTierFast, []fleetEndpoint{{
|
|
ProviderID: "cpu", ModelID: "cpu-model", Tier: capability.FleetTierFast, URL: server.URL,
|
|
Headers: map[string]string{"Authorization": "Bearer fleet-key"},
|
|
}})
|
|
recorder := httptest.NewRecorder()
|
|
handler.ServeHTTP(recorder, httptest.NewRequest(http.MethodGet, "/v1/fleet/readiness", nil))
|
|
if recorder.Code != http.StatusOK || !strings.Contains(recorder.Body.String(), `"ready":true`) {
|
|
t.Fatalf("status=%d body=%s", recorder.Code, recorder.Body.String())
|
|
}
|
|
}
|