87 lines
2.5 KiB
Go
87 lines
2.5 KiB
Go
package llm
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
func mockServer(t *testing.T, response string) *httptest.Server {
|
|
t.Helper()
|
|
return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
assert.Equal(t, "/chat/completions", r.URL.Path)
|
|
assert.Equal(t, "application/json", r.Header.Get("Content-Type"))
|
|
w.Header().Set("Content-Type", "application/json")
|
|
_ = json.NewEncoder(w).Encode(map[string]any{
|
|
"choices": []map[string]any{
|
|
{"message": map[string]any{"role": "assistant", "content": response}},
|
|
},
|
|
})
|
|
}))
|
|
}
|
|
|
|
func TestClient_Complete(t *testing.T) {
|
|
srv := mockServer(t, "hello world")
|
|
defer srv.Close()
|
|
|
|
c := New(srv.URL, "", "test-model", 10*time.Second)
|
|
got, err := c.Complete(context.Background(), "you are helpful", "say hello")
|
|
require.NoError(t, err)
|
|
assert.Equal(t, "hello world", got)
|
|
}
|
|
|
|
func TestClient_ReturnsErrorOnNon200(t *testing.T) {
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
http.Error(w, "overloaded", http.StatusServiceUnavailable)
|
|
}))
|
|
defer srv.Close()
|
|
|
|
c := New(srv.URL, "", "test-model", 10*time.Second)
|
|
_, err := c.Complete(context.Background(), "sys", "user")
|
|
assert.Error(t, err)
|
|
}
|
|
|
|
func TestClient_SendsAuthHeader(t *testing.T) {
|
|
var gotAuth string
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
gotAuth = r.Header.Get("Authorization")
|
|
_ = json.NewEncoder(w).Encode(map[string]any{
|
|
"choices": []map[string]any{{"message": map[string]any{"content": "ok"}}},
|
|
})
|
|
}))
|
|
defer srv.Close()
|
|
|
|
c := New(srv.URL, "my-key", "test-model", 10*time.Second)
|
|
_, err := c.Complete(context.Background(), "sys", "user")
|
|
require.NoError(t, err)
|
|
assert.Equal(t, "Bearer my-key", gotAuth)
|
|
}
|
|
|
|
func TestClient_Retries429(t *testing.T) {
|
|
calls := 0
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
calls++
|
|
if calls == 1 {
|
|
w.Header().Set("Retry-After", "0")
|
|
w.WriteHeader(http.StatusTooManyRequests)
|
|
return
|
|
}
|
|
_ = json.NewEncoder(w).Encode(map[string]any{
|
|
"choices": []map[string]any{{"message": map[string]any{"content": "retried"}}},
|
|
})
|
|
}))
|
|
defer srv.Close()
|
|
|
|
c := New(srv.URL, "", "test-model", 10*time.Second)
|
|
got, err := c.Complete(context.Background(), "sys", "user")
|
|
require.NoError(t, err)
|
|
assert.Equal(t, "retried", got)
|
|
assert.Equal(t, 2, calls)
|
|
}
|