117 lines
3.8 KiB
Go
117 lines
3.8 KiB
Go
package script
|
|
|
|
import (
|
|
"context"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
func TestExecutorRunsFunctionExpression(t *testing.T) {
|
|
out, err := (Executor{}).Execute(context.Background(), Options{
|
|
Script: `(params) => ({ prompt: params.prompt.toUpperCase(), n: 2 })`,
|
|
Args: []any{map[string]any{"prompt": "hello"}},
|
|
ScriptName: "custom_preprocess_script",
|
|
Timeout: time.Second,
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("execute failed: %v", err)
|
|
}
|
|
result := out.(map[string]any)
|
|
if result["prompt"] != "HELLO" || result["n"].(float64) != 2 {
|
|
t.Fatalf("unexpected result: %#v", result)
|
|
}
|
|
}
|
|
|
|
func TestExecutorSelectsPreferredEntry(t *testing.T) {
|
|
out, err := (Executor{}).Execute(context.Background(), Options{
|
|
Script: `
|
|
function helper() { return { wrong: true }; }
|
|
async function submitTask(payload, context) {
|
|
return { status: "submitted", task_id: payload.id, baseURL: context.baseURL };
|
|
}
|
|
`,
|
|
Args: []any{map[string]any{"id": "task-1"}, map[string]any{"baseURL": "https://example.test"}},
|
|
ContextData: map[string]any{"baseURL": "https://example.test"},
|
|
PreferredEntryNames: []string{"submitTask", "submit"},
|
|
ScriptName: "custom_submit_script:video_generate",
|
|
Timeout: time.Second,
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("execute failed: %v", err)
|
|
}
|
|
result := out.(map[string]any)
|
|
if result["task_id"] != "task-1" || result["baseURL"] != "https://example.test" {
|
|
t.Fatalf("unexpected result: %#v", result)
|
|
}
|
|
}
|
|
|
|
func TestExecutorGotJSONHelper(t *testing.T) {
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
if r.Method != http.MethodPost {
|
|
t.Fatalf("unexpected method: %s", r.Method)
|
|
}
|
|
if r.Header.Get("Authorization") != "Bearer test" {
|
|
t.Fatalf("missing authorization header")
|
|
}
|
|
w.Header().Set("Content-Type", "application/json")
|
|
_, _ = w.Write([]byte(`{"status":"success","task_id":"remote-1"}`))
|
|
}))
|
|
defer server.Close()
|
|
|
|
out, err := (Executor{}).Execute(context.Background(), Options{
|
|
Script: `
|
|
async function submitTask(payload, context) {
|
|
return await got.post(context.baseURL + "/tasks", {
|
|
headers: { Authorization: "Bearer " + context.authValues.apiKey },
|
|
json: payload
|
|
}).json();
|
|
}
|
|
`,
|
|
Args: []any{map[string]any{"prompt": "hello"}, map[string]any{"baseURL": server.URL, "authValues": map[string]any{"apiKey": "test"}}},
|
|
ContextData: map[string]any{"baseURL": server.URL, "authValues": map[string]any{"apiKey": "test"}},
|
|
PreferredEntryNames: []string{"submitTask"},
|
|
ScriptName: "custom_submit_script:image_generate",
|
|
Timeout: 2 * time.Second,
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("execute failed: %v", err)
|
|
}
|
|
result := out.(map[string]any)
|
|
if result["task_id"] != "remote-1" {
|
|
t.Fatalf("unexpected result: %#v", result)
|
|
}
|
|
}
|
|
|
|
func TestExecutorTimeout(t *testing.T) {
|
|
_, err := (Executor{}).Execute(context.Background(), Options{
|
|
Script: `async function main() { await new Promise((resolve) => setTimeout(resolve, 200)); return true; }`,
|
|
ScriptName: "custom_poll_script",
|
|
Timeout: 25 * time.Millisecond,
|
|
})
|
|
if err == nil {
|
|
t.Fatal("expected timeout")
|
|
}
|
|
scriptErr, ok := err.(*Error)
|
|
if !ok || scriptErr.Code != "script_timeout" {
|
|
t.Fatalf("expected script_timeout, got %#v", err)
|
|
}
|
|
}
|
|
|
|
func TestExecutorRejectedPromiseMessage(t *testing.T) {
|
|
_, err := (Executor{}).Execute(context.Background(), Options{
|
|
Script: `async function main() { throw new Error("boom"); }`,
|
|
ScriptName: "custom_submit_script",
|
|
Timeout: time.Second,
|
|
})
|
|
if err == nil {
|
|
t.Fatal("expected rejection")
|
|
}
|
|
scriptErr, ok := err.(*Error)
|
|
if !ok || scriptErr.Code != "script_error" || !strings.Contains(scriptErr.Message, "boom") {
|
|
t.Fatalf("expected script_error with boom, got %#v", err)
|
|
}
|
|
}
|