160 lines
3.6 KiB
Go
160 lines
3.6 KiB
Go
package llm
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"log/slog"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"testing"
|
|
)
|
|
|
|
func TestContentPartMarshalJSON(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
part ContentPart
|
|
want string
|
|
}{
|
|
{
|
|
name: "text",
|
|
part: ContentPart{
|
|
Type: ContentTypeText,
|
|
Text: "hello",
|
|
},
|
|
want: `{"text":"hello","type":"text"}`,
|
|
},
|
|
{
|
|
name: "refusal",
|
|
part: ContentPart{
|
|
Type: ContentTypeRefusal,
|
|
Refusal: "cannot comply",
|
|
},
|
|
want: `{"refusal":"cannot comply","type":"refusal"}`,
|
|
},
|
|
{
|
|
name: "image URL",
|
|
part: ContentPart{
|
|
Type: ContentTypeImageURL,
|
|
ImageURL: &ImageURL{
|
|
URL: "https://example.com/image.png",
|
|
},
|
|
},
|
|
want: `{"image_url":{"url":"https://example.com/image.png"},"type":"image_url"}`,
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
got, err := json.Marshal(tt.part)
|
|
if err != nil {
|
|
t.Fatalf("json.Marshal() error = %v", err)
|
|
}
|
|
if string(got) != tt.want {
|
|
t.Errorf("json.Marshal() = %s, want %s", got, tt.want)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestClientReturnsAPIError(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
call func(*Client) error
|
|
wantErr string
|
|
}{
|
|
{
|
|
name: "completions",
|
|
call: func(client *Client) error {
|
|
_, err := client.Completions(t.Context(), CompletionsParams{})
|
|
return err
|
|
},
|
|
wantErr: "completions: model_not_found: model unavailable",
|
|
},
|
|
{
|
|
name: "completions stream",
|
|
call: func(client *Client) error {
|
|
_, err := client.CompletionsStream(
|
|
t.Context(), CompletionsParams{},
|
|
)
|
|
return err
|
|
},
|
|
wantErr: "completions stream: model_not_found: model unavailable",
|
|
},
|
|
{
|
|
name: "list models",
|
|
call: func(client *Client) error {
|
|
_, err := client.ListModels(context.Background())
|
|
return err
|
|
},
|
|
wantErr: "list models: model_not_found: model unavailable",
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
server := httptest.NewServer(http.HandlerFunc(func(
|
|
w http.ResponseWriter, _ *http.Request,
|
|
) {
|
|
w.WriteHeader(http.StatusNotFound)
|
|
fmt.Fprint(w, `{"error":{"code":"model_not_found","message":"model unavailable","type":"invalid_request_error"}}`)
|
|
}))
|
|
defer server.Close()
|
|
|
|
log := slog.New(slog.NewTextHandler(io.Discard, nil))
|
|
client, err := NewClient(Config{URL: server.URL}, log)
|
|
if err != nil {
|
|
t.Fatalf("NewClient() error = %v", err)
|
|
}
|
|
|
|
err = tt.call(client)
|
|
if err == nil {
|
|
t.Fatal("call error = nil, want API error")
|
|
}
|
|
if err.Error() != tt.wantErr {
|
|
t.Errorf("call error = %q, want %q", err, tt.wantErr)
|
|
}
|
|
|
|
var apiErr *APIError
|
|
if !errors.As(err, &apiErr) {
|
|
t.Fatalf("errors.As(%T) = false, want *APIError", err)
|
|
}
|
|
if apiErr.Code != "model_not_found" {
|
|
t.Errorf("APIError.Code = %q, want %q", apiErr.Code, "model_not_found")
|
|
}
|
|
if apiErr.Message != "model unavailable" {
|
|
t.Errorf("APIError.Message = %q, want %q", apiErr.Message, "model unavailable")
|
|
}
|
|
if apiErr.Type != "invalid_request_error" {
|
|
t.Errorf(
|
|
"APIError.Type = %q, want %q",
|
|
apiErr.Type,
|
|
"invalid_request_error",
|
|
)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestMessageUnmarshalJSONReplacesArrayContent(t *testing.T) {
|
|
message := Message{
|
|
ContentParts: []ContentPart{{
|
|
Type: ContentTypeText,
|
|
Text: "old",
|
|
}},
|
|
}
|
|
|
|
err := json.Unmarshal(
|
|
[]byte(`{"content":[{"type":"text","text":"new"}]}`),
|
|
&message,
|
|
)
|
|
if err != nil {
|
|
t.Fatalf("json.Unmarshal() error = %v", err)
|
|
}
|
|
if got := message.Text(); got != "new" {
|
|
t.Errorf("Message.Text() = %q, want %q", got, "new")
|
|
}
|
|
}
|