Skip to content

Commit 40b1862

Browse files
kjlubickkjlubick
committed
Use table-based tests
Signed-off-by: Kaylee Lubick <klubick@nvidia.com>
1 parent a5500b1 commit 40b1862

2 files changed

Lines changed: 94 additions & 152 deletions

File tree

‎services/lmstudio-proxy/failover_test.go‎

Lines changed: 45 additions & 76 deletions
Original file line numberDiff line numberDiff line change
@@ -599,82 +599,51 @@ func TestResolveCandidates_SelfGuard(t *testing.T) {
599599
}
600600
}
601601

602-
// TestHandleHTTP_AnthropicMessages_InferenceRouting proves /v1/messages is
603-
// treated as an inference endpoint: model-based candidate filtering is applied
604-
// and the request body is forwarded unchanged to the matching node.
605-
func TestHandleHTTP_AnthropicMessages_InferenceRouting(t *testing.T) {
606-
var gotBody string
607-
var gotPath string
608-
good := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
609-
b, _ := io.ReadAll(r.Body)
610-
gotBody = string(b)
611-
gotPath = r.URL.Path
612-
w.WriteHeader(http.StatusOK)
613-
_, _ = io.WriteString(w, `{"id":"msg_test","type":"message","role":"assistant","content":[{"type":"text","text":"Hello!"}]}`)
614-
}))
615-
defer good.Close()
616-
617-
wrongModel := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
618-
t.Error("wrong-model node should not receive request")
619-
w.WriteHeader(http.StatusOK)
620-
}))
621-
defer wrongModel.Close()
622-
623-
disc := NewDiscovery()
624-
disc.AddManual(nodeForModel(t, "wrong", wrongModel.URL, "mistral"))
625-
disc.AddManual(nodeForModel(t, "good", good.URL, "claude-3-5-sonnet-20241022"))
626-
p := testProxy(disc, 11434)
627-
p.SetSelected("wrong")
628-
629-
body := `{"model":"claude-3-5-sonnet-20241022","max_tokens":1024,"messages":[{"role":"user","content":"Hello"}]}`
630-
rec := httptest.NewRecorder()
631-
p.handleHTTP(rec, httptest.NewRequest(http.MethodPost, "/v1/messages", strings.NewReader(body)))
632-
633-
if rec.Code != http.StatusOK {
634-
t.Fatalf("status = %d, want 200", rec.Code)
635-
}
636-
if gotBody != body {
637-
t.Errorf("node got body %q, want the original request body", gotBody)
638-
}
639-
if gotPath != "/v1/messages" {
640-
t.Errorf("path = %q, want /v1/messages", gotPath)
602+
// TestHandleHTTP_InferenceRouting proves each request family applies
603+
// model-based candidate filtering and forwards the request unchanged.
604+
func TestHandleHTTP_InferenceRouting(t *testing.T) {
605+
test := func(name, path string) {
606+
t.Run(name, func(t *testing.T) {
607+
var gotBody string
608+
var gotPath string
609+
good := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
610+
b, _ := io.ReadAll(r.Body)
611+
gotBody = string(b)
612+
gotPath = r.URL.Path
613+
w.WriteHeader(http.StatusOK)
614+
}))
615+
defer good.Close()
616+
617+
wrongModel := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
618+
t.Error("wrong-model node should not receive request")
619+
w.WriteHeader(http.StatusOK)
620+
}))
621+
defer wrongModel.Close()
622+
623+
disc := NewDiscovery()
624+
disc.AddManual(nodeForModel(t, "wrong", wrongModel.URL, "different-model"))
625+
disc.AddManual(nodeForModel(t, "good", good.URL, "requested-model"))
626+
p := testProxy(disc, 11434)
627+
p.SetSelected("wrong")
628+
629+
body := `{"model":"requested-model","messages":[]}`
630+
rec := httptest.NewRecorder()
631+
p.handleHTTP(rec, httptest.NewRequest(http.MethodPost, path, strings.NewReader(body)))
632+
633+
if rec.Code != http.StatusOK {
634+
t.Fatalf("status = %d, want 200", rec.Code)
635+
}
636+
if gotBody != body {
637+
t.Errorf("node got body %q, want %q", gotBody, body)
638+
}
639+
if gotPath != path {
640+
t.Errorf("path = %q, want %q", gotPath, path)
641+
}
642+
})
641643
}
642-
}
643-
644-
// TestHandleHTTP_AnthropicMessages_404Failover proves /v1/messages is treated
645-
// as an inference endpoint for 404 failover: when the first eligible candidate
646-
// returns 404, the request retries to the second candidate and the body is
647-
// preserved.
648-
func TestHandleHTTP_AnthropicMessages_404Failover(t *testing.T) {
649-
missing := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
650-
w.WriteHeader(http.StatusNotFound)
651-
_, _ = io.WriteString(w, `{"error":{"type":"not_found_error","message":"model not found"}}`)
652-
}))
653-
defer missing.Close()
654-
655-
var gotBody string
656-
has := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
657-
b, _ := io.ReadAll(r.Body)
658-
gotBody = string(b)
659-
w.WriteHeader(http.StatusOK)
660-
_, _ = io.WriteString(w, `{"id":"msg_test","type":"message","role":"assistant","content":[{"type":"text","text":"Hello!"}]}`)
661-
}))
662-
defer has.Close()
663-
664-
disc := NewDiscovery()
665-
disc.AddManual(nodeForModel(t, "missing", missing.URL, "claude-3-5-sonnet-20241022"))
666-
disc.AddManual(nodeForModel(t, "has", has.URL, "claude-3-5-sonnet-20241022"))
667-
p := testProxy(disc, 11434)
668-
p.SetSelected("missing")
669644

670-
body := `{"model":"claude-3-5-sonnet-20241022","max_tokens":1024,"messages":[{"role":"user","content":"Hello"}]}`
671-
rec := httptest.NewRecorder()
672-
p.handleHTTP(rec, httptest.NewRequest(http.MethodPost, "/v1/messages", strings.NewReader(body)))
673-
674-
if rec.Code != http.StatusOK {
675-
t.Fatalf("status = %d, want 200 (should fail over past the 404)", rec.Code)
676-
}
677-
if gotBody != body {
678-
t.Errorf("failover node got body %q, want the original request body", gotBody)
679-
}
645+
test("OpenAI chat completions", "/v1/chat/completions")
646+
test("OpenAI completions", "/v1/completions")
647+
test("OpenAI embeddings", "/v1/embeddings")
648+
test("Anthropic messages", "/v1/messages")
680649
}

‎services/ollama-proxy/failover_test.go‎

Lines changed: 49 additions & 76 deletions
Original file line numberDiff line numberDiff line change
@@ -643,82 +643,55 @@ func TestResolveCandidates_SelfGuard(t *testing.T) {
643643
}
644644
}
645645

646-
// TestHandleHTTP_AnthropicMessages_InferenceRouting proves /v1/messages is
647-
// treated as an inference endpoint: model-based candidate filtering is applied
648-
// and the request body is forwarded unchanged to the matching node.
649-
func TestHandleHTTP_AnthropicMessages_InferenceRouting(t *testing.T) {
650-
var gotBody string
651-
var gotPath string
652-
good := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
653-
b, _ := io.ReadAll(r.Body)
654-
gotBody = string(b)
655-
gotPath = r.URL.Path
656-
w.WriteHeader(http.StatusOK)
657-
_, _ = io.WriteString(w, `{"id":"msg_test","type":"message","role":"assistant","content":[{"type":"text","text":"Hello!"}]}`)
658-
}))
659-
defer good.Close()
660-
661-
wrongModel := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
662-
t.Error("wrong-model node should not receive request")
663-
w.WriteHeader(http.StatusOK)
664-
}))
665-
defer wrongModel.Close()
666-
667-
disc := NewDiscovery()
668-
disc.AddManual(nodeForModel(t, "wrong", wrongModel.URL, "mistral"))
669-
disc.AddManual(nodeForModel(t, "good", good.URL, "claude-3-5-sonnet-20241022"))
670-
p := testProxy(disc, 11434)
671-
p.SetSelected("wrong")
672-
673-
body := `{"model":"claude-3-5-sonnet-20241022","max_tokens":1024,"messages":[{"role":"user","content":"Hello"}]}`
674-
rec := httptest.NewRecorder()
675-
p.handleHTTP(rec, httptest.NewRequest(http.MethodPost, "/v1/messages", strings.NewReader(body)))
676-
677-
if rec.Code != http.StatusOK {
678-
t.Fatalf("status = %d, want 200", rec.Code)
679-
}
680-
if gotBody != body {
681-
t.Errorf("node got body %q, want the original request body", gotBody)
682-
}
683-
if gotPath != "/v1/messages" {
684-
t.Errorf("path = %q, want /v1/messages", gotPath)
646+
// TestHandleHTTP_InferenceRouting proves each request family applies
647+
// model-based candidate filtering and forwards the request unchanged.
648+
func TestHandleHTTP_InferenceRouting(t *testing.T) {
649+
test := func(name, path string) {
650+
t.Run(name, func(t *testing.T) {
651+
var gotBody string
652+
var gotPath string
653+
good := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
654+
b, _ := io.ReadAll(r.Body)
655+
gotBody = string(b)
656+
gotPath = r.URL.Path
657+
w.WriteHeader(http.StatusOK)
658+
}))
659+
defer good.Close()
660+
661+
wrongModel := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
662+
t.Error("wrong-model node should not receive request")
663+
w.WriteHeader(http.StatusOK)
664+
}))
665+
defer wrongModel.Close()
666+
667+
disc := NewDiscovery()
668+
disc.AddManual(nodeForModel(t, "wrong", wrongModel.URL, "different-model"))
669+
disc.AddManual(nodeForModel(t, "good", good.URL, "requested-model"))
670+
p := testProxy(disc, 11434)
671+
p.SetSelected("wrong")
672+
673+
body := `{"model":"requested-model","messages":[]}`
674+
rec := httptest.NewRecorder()
675+
p.handleHTTP(rec, httptest.NewRequest(http.MethodPost, path, strings.NewReader(body)))
676+
677+
if rec.Code != http.StatusOK {
678+
t.Fatalf("status = %d, want 200", rec.Code)
679+
}
680+
if gotBody != body {
681+
t.Errorf("node got body %q, want %q", gotBody, body)
682+
}
683+
if gotPath != path {
684+
t.Errorf("path = %q, want %q", gotPath, path)
685+
}
686+
})
685687
}
686-
}
687-
688-
// TestHandleHTTP_AnthropicMessages_404Failover proves /v1/messages is treated
689-
// as an inference endpoint for 404 failover: when the first eligible candidate
690-
// returns 404, the request retries to the second candidate and the body is
691-
// preserved.
692-
func TestHandleHTTP_AnthropicMessages_404Failover(t *testing.T) {
693-
missing := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
694-
w.WriteHeader(http.StatusNotFound)
695-
_, _ = io.WriteString(w, `{"error":{"type":"not_found_error","message":"model not found"}}`)
696-
}))
697-
defer missing.Close()
698-
699-
var gotBody string
700-
has := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
701-
b, _ := io.ReadAll(r.Body)
702-
gotBody = string(b)
703-
w.WriteHeader(http.StatusOK)
704-
_, _ = io.WriteString(w, `{"id":"msg_test","type":"message","role":"assistant","content":[{"type":"text","text":"Hello!"}]}`)
705-
}))
706-
defer has.Close()
707-
708-
disc := NewDiscovery()
709-
disc.AddManual(nodeForModel(t, "missing", missing.URL, "claude-3-5-sonnet-20241022"))
710-
disc.AddManual(nodeForModel(t, "has", has.URL, "claude-3-5-sonnet-20241022"))
711-
p := testProxy(disc, 11434)
712-
p.SetSelected("missing")
713688

714-
body := `{"model":"claude-3-5-sonnet-20241022","max_tokens":1024,"messages":[{"role":"user","content":"Hello"}]}`
715-
rec := httptest.NewRecorder()
716-
p.handleHTTP(rec, httptest.NewRequest(http.MethodPost, "/v1/messages", strings.NewReader(body)))
717-
718-
if rec.Code != http.StatusOK {
719-
t.Fatalf("status = %d, want 200 (should fail over past the 404)", rec.Code)
720-
}
721-
if gotBody != body {
722-
t.Errorf("failover node got body %q, want the original request body", gotBody)
723-
}
689+
test("native Ollama generate", "/api/generate")
690+
test("native Ollama chat", "/api/chat")
691+
test("native Ollama embeddings", "/api/embeddings")
692+
test("native Ollama embed", "/api/embed")
693+
test("OpenAI chat completions", "/v1/chat/completions")
694+
test("OpenAI completions", "/v1/completions")
695+
test("OpenAI embeddings", "/v1/embeddings")
696+
test("Anthropic messages", "/v1/messages")
724697
}

0 commit comments

Comments
 (0)