ce7cb702218693d08ea6d83464d2457c9250027c

services/vertex-batch/cloud_test.go

@@ -0,0 +1,116 @@
+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)
+			}
+		})
+	}
+}