74 lines
2.5 KiB
Go
74 lines
2.5 KiB
Go
package openrouter
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"io"
|
|
"net/http"
|
|
"strings"
|
|
"testing"
|
|
)
|
|
|
|
type roundTripFunc func(*http.Request) (*http.Response, error)
|
|
|
|
func (f roundTripFunc) RoundTrip(r *http.Request) (*http.Response, error) { return f(r) }
|
|
func response(status int, contentType, body string) *http.Response {
|
|
return &http.Response{StatusCode: status, Header: http.Header{"Content-Type": []string{contentType}}, Body: io.NopCloser(strings.NewReader(body))}
|
|
}
|
|
|
|
func TestPromptMultimodalRequest(t *testing.T) {
|
|
var body map[string]any
|
|
httpClient := &http.Client{Transport: roundTripFunc(func(r *http.Request) (*http.Response, error) {
|
|
if r.URL.Path != "/chat/completions" {
|
|
t.Errorf("path %s", r.URL.Path)
|
|
}
|
|
if r.Header.Get("Authorization") != "Bearer secret" {
|
|
t.Error("missing auth")
|
|
}
|
|
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
return response(200, "application/json", `{"choices":[{"message":{"content":"{\"video_prompt\":\"gentle motion\"}"}}]}`), nil
|
|
})}
|
|
c := Client{BaseURL: "https://openrouter.test", APIKey: "secret", HTTP: httpClient}
|
|
got, err := c.Prompt(context.Background(), "model", "student text", "data:image/png;base64,abc", "system")
|
|
if err != nil || got != "gentle motion" {
|
|
t.Fatalf("got %q %v", got, err)
|
|
}
|
|
encoded, _ := json.Marshal(body)
|
|
for _, want := range []string{"student text", "data:image/png;base64,abc", "json_schema"} {
|
|
if !strings.Contains(string(encoded), want) {
|
|
t.Errorf("request missing %q", want)
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestVideoFlow(t *testing.T) {
|
|
httpClient := &http.Client{Transport: roundTripFunc(func(r *http.Request) (*http.Response, error) {
|
|
switch {
|
|
case r.Method == http.MethodPost && r.URL.Path == "/videos":
|
|
return response(202, "application/json", `{"id":"job1","status":"pending"}`), nil
|
|
case r.URL.Path == "/videos/job1":
|
|
return response(200, "application/json", `{"id":"job1","status":"completed"}`), nil
|
|
case r.URL.Path == "/videos/job1/content":
|
|
return response(200, "video/mp4", "video"), nil
|
|
default:
|
|
return response(404, "text/plain", "missing"), nil
|
|
}
|
|
})}
|
|
c := Client{BaseURL: "https://openrouter.test", APIKey: "x", HTTP: httpClient}
|
|
job, err := c.Submit(context.Background(), VideoRequest{Model: "m"})
|
|
if err != nil || job.ID != "job1" {
|
|
t.Fatal(job, err)
|
|
}
|
|
job, err = c.Poll(context.Background(), job.ID)
|
|
if err != nil || job.Status != "completed" {
|
|
t.Fatal(job, err)
|
|
}
|
|
resp, err := c.Content(context.Background(), job.ID)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
resp.Body.Close()
|
|
}
|