116 lines
4.1 KiB
Go
116 lines
4.1 KiB
Go
package main
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"sync"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
func TestCloud(t *testing.T) {
|
|
for _, test := range []struct {
|
|
name, state, success string
|
|
cancel bool
|
|
wantError bool
|
|
}{
|
|
{"success", "JOB_STATE_SUCCEEDED", "1", false, false},
|
|
{"failed row", "JOB_STATE_SUCCEEDED", "0", false, true},
|
|
{"failed job", "JOB_STATE_FAILED", "0", false, true},
|
|
{"cancel", "JOB_STATE_PENDING", "0", true, true},
|
|
} {
|
|
t.Run(test.name, func(t *testing.T) {
|
|
var mu sync.Mutex
|
|
objects := map[string]bool{}
|
|
input, prefix := "", ""
|
|
cancelled := false
|
|
submitted := make(chan struct{})
|
|
response := json.RawMessage(`{"candidates":[{"content":{"parts":[{"text":"answer"}]},"finishReason":"STOP"}],"usageMetadata":{"totalTokenCount":16}}`)
|
|
native := json.RawMessage(`{"contents":[{"role":"user","parts":[{"text":"synthetic"},{"inlineData":{"mimeType":"image/png","data":"aGVsbG8="}}]}],"generationConfig":{"temperature":0.2}}`)
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
mu.Lock()
|
|
defer mu.Unlock()
|
|
w.Header().Set("Content-Type", "application/json")
|
|
switch {
|
|
case strings.HasPrefix(r.URL.Path, "/upload/"):
|
|
input = r.URL.Query().Get("name")
|
|
objects[input] = true
|
|
var row struct {
|
|
Request json.RawMessage `json:"request"`
|
|
}
|
|
if err := json.NewDecoder(r.Body).Decode(&row); err != nil {
|
|
t.Error(err)
|
|
}
|
|
if string(row.Request) != string(native) {
|
|
t.Errorf("native body changed: %s", row.Request)
|
|
}
|
|
_, _ = fmt.Fprint(w, `{}`)
|
|
case strings.HasSuffix(r.URL.Path, "/batchPredictionJobs"):
|
|
var body struct {
|
|
Model string `json:"model"`
|
|
Output struct {
|
|
GCS struct {
|
|
Prefix string `json:"outputUriPrefix"`
|
|
} `json:"gcsDestination"`
|
|
} `json:"outputConfig"`
|
|
}
|
|
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
|
|
t.Error(err)
|
|
}
|
|
if body.Model != "publishers/google/models/gemini-test" {
|
|
t.Error(body.Model)
|
|
}
|
|
prefix = strings.TrimPrefix(body.Output.GCS.Prefix, "gs://synthetic/")
|
|
_, _ = fmt.Fprint(w, `{"name":"projects/synthetic/locations/global/batchPredictionJobs/1","state":"JOB_STATE_PENDING"}`)
|
|
close(submitted)
|
|
case strings.HasSuffix(r.URL.Path, ":cancel"):
|
|
cancelled = true
|
|
_, _ = fmt.Fprint(w, `{}`)
|
|
case strings.Contains(r.URL.Path, "/batchPredictionJobs/"):
|
|
objects[prefix+"predictions.jsonl"] = true
|
|
_, _ = fmt.Fprintf(w, `{"name":"projects/synthetic/locations/global/batchPredictionJobs/1","state":%q,"completionStats":{"successfulCount":%q}}`, test.state, test.success)
|
|
case r.Method == "DELETE":
|
|
name := strings.SplitN(r.URL.Path, "/o/", 2)[1]
|
|
delete(objects, name)
|
|
w.WriteHeader(204)
|
|
case r.URL.Query().Get("alt") == "media":
|
|
_, _ = fmt.Fprintf(w, `{"status":"","response":%s}`, response)
|
|
case strings.HasSuffix(r.URL.Path, "/o"):
|
|
var items []map[string]string
|
|
for name := range objects {
|
|
if strings.HasPrefix(name, r.URL.Query().Get("prefix")) {
|
|
items = append(items, map[string]string{"name": name})
|
|
}
|
|
}
|
|
_ = json.NewEncoder(w).Encode(map[string]any{"items": items})
|
|
default:
|
|
t.Error(r.Method, r.URL.String())
|
|
w.WriteHeader(404)
|
|
}
|
|
}))
|
|
defer server.Close()
|
|
client := &cloud{client: server.Client(), vertex: server.URL, storage: server.URL, project: "synthetic", bucket: "synthetic", poll: time.Millisecond}
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
defer cancel()
|
|
if test.cancel {
|
|
go func() { <-submitted; cancel() }()
|
|
}
|
|
result, err := client.run(ctx, "gemini-test", native)
|
|
if (err != nil) != test.wantError {
|
|
t.Fatalf("result %s, error %v", result, err)
|
|
}
|
|
if !test.wantError && string(result) != string(response) {
|
|
t.Errorf("result changed: %s", result)
|
|
}
|
|
mu.Lock()
|
|
defer mu.Unlock()
|
|
if len(objects) != 0 || cancelled != test.cancel {
|
|
t.Errorf("cleanup objects=%v cancelled=%v", objects, cancelled)
|
|
}
|
|
})
|
|
}
|
|
}
|