package vision import ( "context" "encoding/base64" "encoding/json" "net/http" "net/http/httptest" "testing" "time" ) // TestEmbedReturnsVectorsAndSkipsBroken 验证正常路径:批量编码返回向量, // 服务端跳过的图片下标原样带回,向量维度与配置一致。 func TestEmbedReturnsVectorsAndSkipsBroken(t *testing.T) { srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if r.URL.Path != "/embed" { t.Errorf("路径不对: %s", r.URL.Path) } _ = json.NewEncoder(w).Encode(map[string]any{ "dim": 768, "vectors": [][]float32{make([]float32, 768), make([]float32, 768)}, "indexes": []int{0, 2}, "skipped": []int{1}, "device": "cpu", }) })) defer srv.Close() c := New(srv.URL, 768, 2*time.Second, 3) vecs, skipped, err := c.Embed(context.Background(), [][]byte{[]byte("a"), []byte("b"), []byte("c")}) if err != nil { t.Fatalf("Embed 返回错误: %v", err) } if len(vecs) != 2 { t.Fatalf("应有 2 个向量,得到 %d", len(vecs)) } if len(skipped) != 1 || skipped[0] != 1 { t.Fatalf("skipped 应为 [1],得到 %v", skipped) } if len(vecs[0]) != 768 { t.Fatalf("向量维度应为 768,得到 %d", len(vecs[0])) } } // TestEmbedDisabledSkipsCall 验证禁用后完全不发起 HTTP 请求(视觉服务未启用时的零开销路径)。 func TestEmbedDisabledSkipsCall(t *testing.T) { called := false srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { called = true })) defer srv.Close() c := New(srv.URL, 768, time.Second, 3) c.Disable() vecs, skipped, err := c.Embed(context.Background(), [][]byte{[]byte("x")}) if err != ErrUnavailable { t.Fatalf("禁用时应返回 ErrUnavailable,实际 %v", err) } if vecs != nil || skipped != nil { t.Fatalf("禁用时应返回 nil,实际 %v %v", vecs, skipped) } if called { t.Fatal("禁用时不应发起 HTTP 请求") } } // TestEmbedDimensionMismatchIsError 验证维度不符必须报错: // 若放过,会把错误维度的向量写进 image_embeddings 造成脏数据。 func TestEmbedDimensionMismatchIsError(t *testing.T) { srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { _ = json.NewEncoder(w).Encode(map[string]any{ "dim": 512, "vectors": [][]float32{make([]float32, 512)}, "indexes": []int{0}, "skipped": []int{}, }) })) defer srv.Close() c := New(srv.URL, 768, time.Second, 3) if _, _, err := c.Embed(context.Background(), [][]byte{[]byte("x")}); err == nil { t.Fatal("维度不符应返回错误") } } // TestEmbedTripBreakerThenCooldownSkips 验证熔断:连续失败达阈值后进入冷却, // 冷却期内直接返回错误且不再发起请求(避免每张图都卡在超时上)。 func TestEmbedTripBreakerThenCooldownSkips(t *testing.T) { hits := 0 srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { hits++ w.WriteHeader(http.StatusInternalServerError) })) defer srv.Close() c := New(srv.URL, 768, time.Second, 3) for i := 0; i < 3; i++ { if _, _, err := c.Embed(context.Background(), [][]byte{[]byte("x")}); err == nil { t.Fatalf("第 %d 次调用应报错", i+1) } } if c.Open() { t.Fatal("熔断后 Open 应为 false") } before := hits if _, _, err := c.Embed(context.Background(), [][]byte{[]byte("x")}); err == nil { t.Fatalf("熔断期应直接返回错误: %v", err) } if hits != before { t.Fatalf("熔断期间不应再发起请求,hits 从 %d 变为 %d", before, hits) } } // TestEmbedSendsBase64Payload 验证请求体格式(base64 编码的图片数组), // 这是与服务端约定的契约,写错会导致服务端解码失败。 func TestEmbedSendsBase64Payload(t *testing.T) { var got struct { Images []string `json:"images"` } srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { _ = json.NewDecoder(r.Body).Decode(&got) _ = json.NewEncoder(w).Encode(map[string]any{ "dim": 768, "vectors": [][]float32{make([]float32, 768)}, "indexes": []int{0}, "skipped": []int{}, }) })) defer srv.Close() c := New(srv.URL, 768, time.Second, 3) if _, _, err := c.Embed(context.Background(), [][]byte{[]byte("hello")}); err != nil { t.Fatalf("Embed 返回错误: %v", err) } want := base64.StdEncoding.EncodeToString([]byte("hello")) if len(got.Images) != 1 || got.Images[0] != want { t.Fatalf("请求体 base64 不正确: %v", got.Images) } } // TestEmbedEmptyInputIsNoop 验证空输入不发起请求、不报错。 func TestEmbedEmptyInputIsNoop(t *testing.T) { called := false srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { called = true })) defer srv.Close() c := New(srv.URL, 768, time.Second, 3) vecs, skipped, err := c.Embed(context.Background(), nil) if err != nil || vecs != nil || skipped != nil { t.Fatalf("空输入应返回 nil,nil,nil,实际 %v %v %v", vecs, skipped, err) } if called { t.Fatal("空输入不应发起请求") } } // TestEmbedBatchTooLarge 验证单批上限,超限直接拒绝(避免把一张大图撑爆服务端)。 func TestEmbedBatchTooLarge(t *testing.T) { c := New("http://127.0.0.1:1", 768, time.Second, 3) imgs := make([][]byte, MaxBatch+1) for i := range imgs { imgs[i] = []byte("x") } if _, _, err := c.Embed(context.Background(), imgs); err != ErrBatchTooLarge { t.Fatalf("超上限应返回 ErrBatchTooLarge,实际 %v", err) } }