diff --git a/pinecone/client.go b/pinecone/client.go index 97e8e91..b34e017 100644 --- a/pinecone/client.go +++ b/pinecone/client.go @@ -386,12 +386,13 @@ func toCollection(cm *control.CollectionModel) *Collection { if cm == nil { return nil } + return &Collection{ Name: cm.Name, - Size: cm.Size, + Size: derefOrDefault(cm.Size, 0), Status: CollectionStatus(cm.Status), - Dimension: cm.Dimension, - VectorCount: cm.VectorCount, + Dimension: derefOrDefault(cm.Dimension, 0), + VectorCount: derefOrDefault(cm.VectorCount, 0), Environment: cm.Environment, } } @@ -413,6 +414,13 @@ func minOne(x int32) int32 { return x } +func derefOrDefault[T any](ptr *T, defaultValue T) T { + if ptr == nil { + return defaultValue + } + return *ptr +} + func buildClientOptions(in NewClientParams) ([]control.ClientOption, error) { clientOptions := []control.ClientOption{} hasAuthorizationHeader := false diff --git a/pinecone/index_connection.go b/pinecone/index_connection.go index a2daf3a..866d557 100644 --- a/pinecone/index_connection.go +++ b/pinecone/index_connection.go @@ -14,18 +14,18 @@ import ( ) type IndexConnection struct { - Namespace string - apiKey string + Namespace string + apiKey string additionalMetadata map[string]string - dataClient *data.VectorServiceClient - grpcConn *grpc.ClientConn + dataClient *data.VectorServiceClient + grpcConn *grpc.ClientConn } type newIndexParameters struct { - apiKey string - host string - namespace string - sourceTag string + apiKey string + host string + namespace string + sourceTag string additionalMetadata map[string]string } @@ -75,8 +75,8 @@ func (idx *IndexConnection) UpsertVectors(ctx context.Context, in []*Vector) (ui } type FetchVectorsResponse struct { - Vectors map[string]*Vector - Usage *Usage + Vectors map[string]*Vector `json:"vectors,omitempty"` + Usage *Usage `json:"usage,omitempty"` } func (idx *IndexConnection) FetchVectors(ctx context.Context, ids []string) (*FetchVectorsResponse, error) { @@ -94,6 +94,7 @@ func (idx *IndexConnection) FetchVectors(ctx context.Context, ids []string) (*Fe for id, vector := range res.Vectors { vectors[id] = toVector(vector) } + fmt.Printf("VECTORS: %+v\n", vectors) return &FetchVectorsResponse{ Vectors: vectors, @@ -108,9 +109,9 @@ type ListVectorsRequest struct { } type ListVectorsResponse struct { - VectorIds []*string - Usage *Usage - NextPaginationToken *string + VectorIds []*string `json:"vector_ids,omitempty"` + Usage *Usage `json:"usage,omitempty"` + NextPaginationToken *string `json:"next_pagination_token,omitempty"` } func (idx *IndexConnection) ListVectors(ctx context.Context, in *ListVectorsRequest) (*ListVectorsResponse, error) { @@ -132,7 +133,7 @@ func (idx *IndexConnection) ListVectors(ctx context.Context, in *ListVectorsRequ return &ListVectorsResponse{ VectorIds: vectorIds, - Usage: &Usage{ReadUnits: res.Usage.ReadUnits}, + Usage: &Usage{ReadUnits: derefOrDefault(res.Usage.ReadUnits, 0)}, NextPaginationToken: toPaginationToken(res.Pagination), }, nil } @@ -147,8 +148,8 @@ type QueryByVectorValuesRequest struct { } type QueryVectorsResponse struct { - Matches []*ScoredVector - Usage *Usage + Matches []*ScoredVector `json:"matches,omitempty"` + Usage *Usage `json:"usage,omitempty"` } func (idx *IndexConnection) QueryByVectorValues(ctx context.Context, in *QueryByVectorValuesRequest) (*QueryVectorsResponse, error) { @@ -236,10 +237,10 @@ func (idx *IndexConnection) UpdateVector(ctx context.Context, in *UpdateVectorRe } type DescribeIndexStatsResponse struct { - Dimension uint32 - IndexFullness float32 - TotalVectorCount uint32 - Namespaces map[string]*NamespaceSummary + Dimension uint32 `json:"dimension"` + IndexFullness float32 `json:"index_fullness"` + TotalVectorCount uint32 `json:"total_vector_count"` + Namespaces map[string]*NamespaceSummary `json:"namespaces,omitempty"` } func (idx *IndexConnection) DescribeIndexStats(ctx context.Context) (*DescribeIndexStatsResponse, error) { @@ -335,7 +336,7 @@ func toUsage(u *data.Usage) *Usage { return nil } return &Usage{ - ReadUnits: u.ReadUnits, + ReadUnits: derefOrDefault(u.ReadUnits, 0), } } @@ -372,7 +373,7 @@ func (idx *IndexConnection) akCtx(ctx context.Context) context.Context { newMetadata := []string{} newMetadata = append(newMetadata, "api-key", idx.apiKey) - for key, value := range idx.additionalMetadata{ + for key, value := range idx.additionalMetadata { newMetadata = append(newMetadata, key, value) } diff --git a/pinecone/index_connection_test.go b/pinecone/index_connection_test.go index fd709d6..7589535 100644 --- a/pinecone/index_connection_test.go +++ b/pinecone/index_connection_test.go @@ -2,6 +2,7 @@ package pinecone import ( "context" + "encoding/json" "fmt" "os" "reflect" @@ -10,6 +11,7 @@ import ( "github.com/google/uuid" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/suite" + "google.golang.org/protobuf/types/known/structpb" ) type IndexConnectionTests struct { @@ -17,9 +19,9 @@ type IndexConnectionTests struct { host string dimension int32 apiKey string + indexType string idxConn *IndexConnection sourceTag string - metadata map[string]string idxConnSourceTag *IndexConnection vectorIds []string } @@ -43,6 +45,7 @@ func TestIndexConnection(t *testing.T) { } podTestSuite := new(IndexConnectionTests) + podTestSuite.indexType = "pod" podTestSuite.host = podIdx.Host podTestSuite.dimension = podIdx.Dimension podTestSuite.apiKey = apiKey @@ -56,6 +59,7 @@ func TestIndexConnection(t *testing.T) { } serverlessTestSuite := new(IndexConnectionTests) + serverlessTestSuite.indexType = "serverless" serverlessTestSuite.host = serverlessIdx.Host serverlessTestSuite.dimension = serverlessIdx.Dimension serverlessTestSuite.apiKey = apiKey @@ -155,9 +159,9 @@ func (ts *IndexConnectionTests) TestNewIndexConnectionAdditionalMetadata() { if idxConn.Namespace != namespace { ts.FailNow(fmt.Sprintf("Expected idxConn to have namespace '%s', but got '%s'", namespace, idxConn.Namespace)) } - if !reflect.DeepEqual(idxConn.additionalMetadata, additionalMetadata) { + if !reflect.DeepEqual(idxConn.additionalMetadata, additionalMetadata) { ts.FailNow(fmt.Sprintf("Expected idxConn to have additionalMetadata '%+v', but got '%+v'", additionalMetadata, idxConn.additionalMetadata)) - } + } if idxConn.dataClient == nil { ts.FailNow("Expected idxConn to have non-nil dataClient") } @@ -247,9 +251,23 @@ func (ts *IndexConnectionTests) TestDeleteVectorsById() { } func (ts *IndexConnectionTests) TestDeleteVectorsByFilter() { + metadataFilter := map[string]interface{}{ + "genre": "classical", + } + filter, err := structpb.NewStruct(metadataFilter) + if err != nil { + ts.FailNow(fmt.Sprintf("Failed to create metadata filter: %v", err)) + } + ctx := context.Background() - err := ts.idxConn.DeleteVectorsByFilter(ctx, &Filter{}) - assert.NoError(ts.T(), err) + err = ts.idxConn.DeleteVectorsByFilter(ctx, filter) + + if ts.indexType == "serverless" { + assert.Error(ts.T(), err) + assert.Containsf(ts.T(), err.Error(), "Serverless and Starter indexes do not support deleting with metadata filtering", "Expected error message to contain 'Serverless and Starter indexes do not support deleting with metadata filtering'") + } else { + assert.NoError(ts.T(), err) + } ts.loadData() //reload deleted data } @@ -286,6 +304,192 @@ func (ts *IndexConnectionTests) TestListVectors() { assert.NotNil(ts.T(), res) } +func TestMarshalFetchVectorsResponse(t *testing.T) { + tests := []struct { + name string + input FetchVectorsResponse + want string + }{ + { + name: "All fields present", + input: FetchVectorsResponse{ + Vectors: map[string]*Vector{ + "vec-1": {Id: "vec-1", Values: []float32{0.01, 0.01, 0.01}}, + "vec-2": {Id: "vec-2", Values: []float32{0.02, 0.02, 0.02}}, + }, + Usage: &Usage{ReadUnits: 5}, + }, + want: `{"vectors":{"vec-1":{"id":"vec-1","values":[0.01,0.01,0.01]},"vec-2":{"id":"vec-2","values":[0.02,0.02,0.02]}},"usage":{"read_units":5}}`, + }, + { + name: "Fields omitted", + input: FetchVectorsResponse{}, + want: `{}`, + }, + { + name: "Fields empty", + input: FetchVectorsResponse{ + Vectors: nil, + Usage: nil, + }, + want: `{}`, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + bytes, err := json.Marshal(tt.input) + if err != nil { + t.Fatalf("Failed to marshal FetchVectorsResponse: %v", err) + } + + if got := string(bytes); got != tt.want { + t.Errorf("Marshal FetchVectorsResponse got = %s, want = %s", got, tt.want) + } + }) + } +} + +func TestMarshalListVectorsResponse(t *testing.T) { + vectorId1 := "vec-1" + vectorId2 := "vec-2" + paginationToken := "next-token" + tests := []struct { + name string + input ListVectorsResponse + want string + }{ + { + name: "All fields present", + input: ListVectorsResponse{ + VectorIds: []*string{&vectorId1, &vectorId2}, + Usage: &Usage{ReadUnits: 5}, + NextPaginationToken: &paginationToken, + }, + want: `{"vector_ids":["vec-1","vec-2"],"usage":{"read_units":5},"next_pagination_token":"next-token"}`, + }, + { + name: "Fields omitted", + input: ListVectorsResponse{}, + want: `{}`, + }, + { + name: "Fields empty", + input: ListVectorsResponse{ + VectorIds: nil, + Usage: nil, + NextPaginationToken: nil, + }, + want: `{}`, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + bytes, err := json.Marshal(tt.input) + if err != nil { + t.Fatalf("Failed to marshal ListVectorsResponse: %v", err) + } + + if got := string(bytes); got != tt.want { + t.Errorf("Marshal ListVectorsResponse got = %s, want = %s", got, tt.want) + } + }) + } +} + +func TestMarshalQueryVectorsResponse(t *testing.T) { + tests := []struct { + name string + input QueryVectorsResponse + want string + }{ + { + name: "All fields present", + input: QueryVectorsResponse{ + Matches: []*ScoredVector{ + {Vector: &Vector{Id: "vec-1", Values: []float32{0.01, 0.01, 0.01}}, Score: 0.1}, + {Vector: &Vector{Id: "vec-2", Values: []float32{0.02, 0.02, 0.02}}, Score: 0.2}, + }, + Usage: &Usage{ReadUnits: 5}, + }, + want: `{"matches":[{"vector":{"id":"vec-1","values":[0.01,0.01,0.01]},"score":0.1},{"vector":{"id":"vec-2","values":[0.02,0.02,0.02]},"score":0.2}],"usage":{"read_units":5}}`, + }, + { + name: "Fields omitted", + input: QueryVectorsResponse{}, + want: `{}`, + }, + { + name: "Fields empty", + input: QueryVectorsResponse{Matches: nil, Usage: nil}, + want: `{}`, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + bytes, err := json.Marshal(tt.input) + if err != nil { + t.Fatalf("Failed to marshal QueryVectorsResponse: %v", err) + } + + if got := string(bytes); got != tt.want { + t.Errorf("Marshal QueryVectorsResponse got = %s, want = %s", got, tt.want) + } + }) + } +} + +func TestMarshalDescribeIndexStatsResponse(t *testing.T) { + tests := []struct { + name string + input DescribeIndexStatsResponse + want string + }{ + { + name: "All fields present", + input: DescribeIndexStatsResponse{ + Dimension: 3, + IndexFullness: 0.5, + TotalVectorCount: 100, + Namespaces: map[string]*NamespaceSummary{ + "namespace-1": {VectorCount: 50}, + }, + }, + want: `{"dimension":3,"index_fullness":0.5,"total_vector_count":100,"namespaces":{"namespace-1":{"vector_count":50}}}`, + }, + { + name: "Fields omitted", + input: DescribeIndexStatsResponse{}, + want: `{"dimension":0,"index_fullness":0,"total_vector_count":0}`, + }, + { + name: "Fields empty", + input: DescribeIndexStatsResponse{ + Dimension: 0, + IndexFullness: 0, + TotalVectorCount: 0, + Namespaces: nil, + }, + want: `{"dimension":0,"index_fullness":0,"total_vector_count":0}`, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + bytes, err := json.Marshal(tt.input) + if err != nil { + t.Fatalf("Failed to marshal DescribeIndexStatsResponse: %v", err) + } + + if got := string(bytes); got != tt.want { + t.Errorf("Marshal DescribeIndexStatsResponse got = %s, want = %s", got, tt.want) + } + }) + } +} + func (ts *IndexConnectionTests) loadData() { vals := []float32{0.01, 0.02, 0.03, 0.04, 0.05} vectors := make([]*Vector, len(vals)) diff --git a/pinecone/models.go b/pinecone/models.go index 3ce895b..4a43525 100644 --- a/pinecone/models.go +++ b/pinecone/models.go @@ -32,31 +32,31 @@ const ( ) type IndexStatus struct { - Ready bool - State IndexStatusState + Ready bool `json:"ready"` + State IndexStatusState `json:"state"` } type IndexSpec struct { - Pod *PodSpec - Serverless *ServerlessSpec + Pod *PodSpec `json:"pod,omitempty"` + Serverless *ServerlessSpec `json:"serverless,omitempty"` } type Index struct { - Name string - Dimension int32 - Host string - Metric IndexMetric - Spec *IndexSpec - Status *IndexStatus + Name string `json:"name"` + Dimension int32 `json:"dimension"` + Host string `json:"host"` + Metric IndexMetric `json:"metric"` + Spec *IndexSpec `json:"spec,omitempty"` + Status *IndexStatus `json:"status,omitempty"` } type Collection struct { - Name string - Size *int64 - Status CollectionStatus - Dimension *int32 - VectorCount *int32 - Environment string + Name string `json:"name"` + Size int64 `json:"size"` + Status CollectionStatus `json:"status"` + Dimension int32 `json:"dimension"` + VectorCount int32 `json:"vector_count"` + Environment string `json:"environment"` } type CollectionStatus string @@ -68,48 +68,48 @@ const ( ) type PodSpecMetadataConfig struct { - Indexed *[]string + Indexed *[]string `json:"indexed,omitempty"` } type PodSpec struct { - Environment string - PodType string - PodCount int32 - Replicas int32 - ShardCount int32 - SourceCollection *string - MetadataConfig *PodSpecMetadataConfig + Environment string `json:"environment"` + PodType string `json:"pod_type"` + PodCount int32 `json:"pod_count"` + Replicas int32 `json:"replicas"` + ShardCount int32 `json:"shard_count"` + SourceCollection *string `json:"source_collection,omitempty"` + MetadataConfig *PodSpecMetadataConfig `json:"metadata_config,omitempty"` } type ServerlessSpec struct { - Cloud Cloud - Region string + Cloud Cloud `json:"cloud"` + Region string `json:"region"` } type Vector struct { - Id string - Values []float32 - SparseValues *SparseValues - Metadata *Metadata + Id string `json:"id"` + Values []float32 `json:"values,omitempty"` + SparseValues *SparseValues `json:"sparse_values,omitempty"` + Metadata *Metadata `json:"metadata,omitempty"` } type ScoredVector struct { - Vector *Vector - Score float32 + Vector *Vector `json:"vector,omitempty"` + Score float32 `json:"score"` } type SparseValues struct { - Indices []uint32 - Values []float32 + Indices []uint32 `json:"indices,omitempty"` + Values []float32 `json:"values,omitempty"` } type NamespaceSummary struct { - VectorCount uint32 + VectorCount uint32 `json:"vector_count"` } type Usage struct { - ReadUnits *uint32 + ReadUnits uint32 `json:"read_units"` } type Filter = structpb.Struct -type Metadata = structpb.Struct \ No newline at end of file +type Metadata = structpb.Struct diff --git a/pinecone/models_test.go b/pinecone/models_test.go new file mode 100644 index 0000000..21aa392 --- /dev/null +++ b/pinecone/models_test.go @@ -0,0 +1,557 @@ +package pinecone + +import ( + "encoding/json" + "testing" + + "google.golang.org/protobuf/types/known/structpb" +) + +func TestMarshalIndexStatus(t *testing.T) { + tests := []struct { + name string + input IndexStatus + want string + }{ + { + name: "All fields present", + input: IndexStatus{Ready: true, State: "Ready"}, + want: `{"ready":true,"state":"Ready"}`, + }, + { + name: "Fields omitted", + input: IndexStatus{}, + want: `{"ready":false,"state":""}`, + }, + { + name: "Fields empty", + input: IndexStatus{Ready: false, State: ""}, + want: `{"ready":false,"state":""}`, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(c *testing.T) { + got, err := json.Marshal(tt.input) + if err != nil { + c.Errorf("Failed to marshal IndexStatus: %v", err) + return + } + if string(got) != tt.want { + c.Errorf("Marshal IndexStatus got = %s, want = %s", string(got), tt.want) + } + }) + } +} + +func TestMarshalServerlessSpec(t *testing.T) { + tests := []struct { + name string + input ServerlessSpec + want string + }{ + { + name: "All fields present", + input: ServerlessSpec{Cloud: "aws", Region: "us-west-"}, + want: `{"cloud":"aws","region":"us-west-"}`, + }, + { + name: "Fields omitted", + input: ServerlessSpec{}, + want: `{"cloud":"","region":""}`, + }, + { + name: "Fields empty", + input: ServerlessSpec{Cloud: "", Region: ""}, + want: `{"cloud":"","region":""}`, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(c *testing.T) { + got, err := json.Marshal(tt.input) + if err != nil { + c.Errorf("Failed to marshal ServerlessSpec: %v", err) + return + } + if string(got) != tt.want { + c.Errorf("Marshal ServerlessSpec got = %s, want = %s", string(got), tt.want) + } + + }) + } +} + +func TestMarshalPodSpec(t *testing.T) { + sourceCollection := "source-collection" + tests := []struct { + name string + input PodSpec + want string + }{ + { + name: "All fields present", + input: PodSpec{ + Environment: "us-west2-gcp", + PodType: "p1.x1", + PodCount: 1, + Replicas: 1, + ShardCount: 1, + SourceCollection: &sourceCollection, + MetadataConfig: &PodSpecMetadataConfig{ + Indexed: &[]string{"genre"}, + }, + }, + want: `{"environment":"us-west2-gcp","pod_type":"p1.x1","pod_count":1,"replicas":1,"shard_count":1,"source_collection":"source-collection","metadata_config":{"indexed":["genre"]}}`, + }, + { + name: "Fields omitted", + input: PodSpec{}, + want: `{"environment":"","pod_type":"","pod_count":0,"replicas":0,"shard_count":0}`, + }, + { + name: "Fields empty", + input: PodSpec{ + Environment: "", + PodType: "", + PodCount: 0, + Replicas: 0, + ShardCount: 0, + SourceCollection: nil, + MetadataConfig: nil, + }, + want: `{"environment":"","pod_type":"","pod_count":0,"replicas":0,"shard_count":0}`, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(c *testing.T) { + got, err := json.Marshal(tt.input) + if err != nil { + c.Errorf("Failed to marshal PodSpec: %v", err) + return + } + if string(got) != tt.want { + c.Errorf("Marshal PodSpec got = %s, want = %s", string(got), tt.want) + } + }) + } +} + +func TestMarshalIndexSpec(t *testing.T) { + sourceCollection := "source-collection" + tests := []struct { + name string + input IndexSpec + want string + }{ + { + name: "Pod spec", + input: IndexSpec{Pod: &PodSpec{ + Environment: "us-west2-gcp", + PodType: "p1.x1", + PodCount: 1, + Replicas: 1, + ShardCount: 1, + SourceCollection: &sourceCollection, + MetadataConfig: &PodSpecMetadataConfig{ + Indexed: &[]string{"genre"}, + }, + }}, + want: `{"pod":{"environment":"us-west2-gcp","pod_type":"p1.x1","pod_count":1,"replicas":1,"shard_count":1,"source_collection":"source-collection","metadata_config":{"indexed":["genre"]}}}`, + }, + { + name: "Serverless spec", + input: IndexSpec{Serverless: &ServerlessSpec{Cloud: "aws", Region: "us-west-"}}, + want: `{"serverless":{"cloud":"aws","region":"us-west-"}}`, + }, + { + name: "Fields omitted", + input: IndexSpec{}, + want: `{}`, + }, + { + name: "Fields empty", + input: IndexSpec{Pod: nil, Serverless: nil}, + want: `{}`, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(c *testing.T) { + got, err := json.Marshal(tt.input) + if err != nil { + c.Errorf("Failed to marshal IndexSpec: %v", err) + return + } + if string(got) != tt.want { + c.Errorf("Marshal IndexSpec got = %s, want = %s", string(got), tt.want) + } + }) + } +} + +func TestMarshalIndex(t *testing.T) { + tests := []struct { + name string + input Index + want string + }{ + { + name: "All fields present", + input: Index{ + Name: "test-index", + Dimension: 128, + Host: "index-host-1.io", + Metric: "cosine", + Spec: &IndexSpec{ + Serverless: &ServerlessSpec{ + Cloud: "aws", + Region: "us-west-2", + }, + }, + Status: &IndexStatus{ + Ready: true, + State: "Ready", + }, + }, + want: `{"name":"test-index","dimension":128,"host":"index-host-1.io","metric":"cosine","spec":{"serverless":{"cloud":"aws","region":"us-west-2"}},"status":{"ready":true,"state":"Ready"}}`, + }, + { + name: "Fields omitted", + input: Index{}, + want: `{"name":"","dimension":0,"host":"","metric":""}`, + }, + { + name: "Fields empty", + input: Index{ + Name: "", + Dimension: 0, + Host: "", + Metric: "", + Spec: nil, + Status: nil, + }, + want: `{"name":"","dimension":0,"host":"","metric":""}`, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(c *testing.T) { + got, err := json.Marshal(tt.input) + if err != nil { + c.Errorf("Failed to marshal Index: %v", err) + return + } + if string(got) != tt.want { + c.Errorf("Marshal Index got = %s, want = %s", string(got), tt.want) + } + }) + } +} + +func TestMarshalCollection(t *testing.T) { + tests := []struct { + name string + input Collection + want string + }{ + { + name: "All fields present", + input: Collection{ + Name: "test-collection", + Size: 15328, + Status: "Ready", + Dimension: 132, + VectorCount: 15000, + Environment: "us-west-2", + }, + want: `{"name":"test-collection","size":15328,"status":"Ready","dimension":132,"vector_count":15000,"environment":"us-west-2"}`, + }, + { + name: "Fields omitted", + input: Collection{}, + want: `{"name":"","size":0,"status":"","dimension":0,"vector_count":0,"environment":""}`, + }, + { + name: "Fields empty", + input: Collection{ + Name: "", + Size: 0, + Status: "", + Dimension: 0, + VectorCount: 0, + Environment: "", + }, + want: `{"name":"","size":0,"status":"","dimension":0,"vector_count":0,"environment":""}`, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(c *testing.T) { + got, err := json.Marshal(tt.input) + if err != nil { + c.Errorf("Failed to marshal Collection: %v", err) + return + } + if string(got) != tt.want { + c.Errorf("Marshal Collection got = %s, want = %s", string(got), tt.want) + } + }) + } +} + +func TestMarshalPodSpecMetadataConfig(t *testing.T) { + tests := []struct { + name string + input PodSpecMetadataConfig + want string + }{ + { + name: "All fields present", + input: PodSpecMetadataConfig{Indexed: &[]string{"genre", "artist"}}, + want: `{"indexed":["genre","artist"]}`, + }, + { + name: "Fields omitted", + input: PodSpecMetadataConfig{}, + want: `{}`, + }, + { + name: "Fields empty", + input: PodSpecMetadataConfig{Indexed: nil}, + want: `{}`, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(c *testing.T) { + got, err := json.Marshal(tt.input) + if err != nil { + c.Errorf("Failed to marshal PodSpecMetadataConfig: %v", err) + return + } + if string(got) != tt.want { + c.Errorf("Marshal PodSpecMetadataConfig got = %s, want = %s", string(got), tt.want) + } + }) + } +} + +func TestMarshalVector(t *testing.T) { + metadata, err := structpb.NewStruct(map[string]interface{}{"genre": "rock"}) + if err != nil { + t.Fatalf("Failed to create metadata: %v", err) + } + + tests := []struct { + name string + input Vector + want string + }{ + { + name: "All fields present", + input: Vector{ + Id: "vector-1", + Values: []float32{0.1, 0.2, 0.3}, + Metadata: metadata, + SparseValues: &SparseValues{ + Indices: []uint32{1, 2, 3}, + Values: []float32{0.1, 0.2, 0.3}, + }, + }, + want: `{"id":"vector-1","values":[0.1,0.2,0.3],"sparse_values":{"indices":[1,2,3],"values":[0.1,0.2,0.3]},"metadata":{"genre":"rock"}}`, + }, + { + name: "Fields omitted", + input: Vector{}, + want: `{"id":""}`, + }, + { + name: "Fields empty", + input: Vector{Id: "", Values: nil, SparseValues: nil, Metadata: nil}, + want: `{"id":""}`, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(c *testing.T) { + got, err := json.Marshal(tt.input) + if err != nil { + c.Errorf("Failed to marshal Vector: %v", err) + return + } + if string(got) != tt.want { + c.Errorf("Marshal Vector got = %s, want = %s", string(got), tt.want) + } + }) + } +} + +func TestMarshalScoredVector(t *testing.T) { + metadata, err := structpb.NewStruct(map[string]interface{}{"genre": "rock"}) + if err != nil { + t.Fatalf("Failed to create metadata: %v", err) + } + + tests := []struct { + name string + input ScoredVector + want string + }{ + { + name: "All fields present", + input: ScoredVector{ + Vector: &Vector{ + Id: "vector-1", + Values: []float32{0.1, 0.2, 0.3}, + Metadata: metadata, + SparseValues: &SparseValues{ + Indices: []uint32{1, 2, 3}, + Values: []float32{0.1, 0.2, 0.3}, + }, + }, + Score: 0.9, + }, + want: `{"vector":{"id":"vector-1","values":[0.1,0.2,0.3],"sparse_values":{"indices":[1,2,3],"values":[0.1,0.2,0.3]},"metadata":{"genre":"rock"}},"score":0.9}`, + }, + { + name: "Fields omitted", + input: ScoredVector{}, + want: `{"score":0}`, + }, + { + name: "Fields empty", + input: ScoredVector{Vector: nil, Score: 0}, + want: `{"score":0}`, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(c *testing.T) { + got, err := json.Marshal(tt.input) + if err != nil { + c.Errorf("Failed to marshal ScoredVector: %v", err) + return + } + if string(got) != tt.want { + c.Errorf("Marshal ScoredVector got = %s, want = %s", string(got), tt.want) + } + }) + } +} + +func TestMarshalSparseValues(t *testing.T) { + tests := []struct { + name string + input SparseValues + want string + }{ + { + name: "All fields present", + input: SparseValues{ + Indices: []uint32{1, 2, 3}, + Values: []float32{0.1, 0.2, 0.3}, + }, + want: `{"indices":[1,2,3],"values":[0.1,0.2,0.3]}`, + }, + { + name: "Fields omitted", + input: SparseValues{}, + want: `{}`, + }, + { + name: "Fields empty", + input: SparseValues{Indices: nil, Values: nil}, + want: `{}`, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(c *testing.T) { + got, err := json.Marshal(tt.input) + if err != nil { + c.Errorf("Failed to marshal SparseValues: %v", err) + return + } + if string(got) != tt.want { + c.Errorf("Marshal SparseValues got = %s, want = %s", string(got), tt.want) + } + }) + } +} + +func TestMarshalNamespaceSummary(t *testing.T) { + tests := []struct { + name string + input NamespaceSummary + want string + }{ + { + name: "All fields present", + input: NamespaceSummary{VectorCount: 15000}, + want: `{"vector_count":15000}`, + }, + { + name: "Fields omitted", + input: NamespaceSummary{}, + want: `{"vector_count":0}`, + }, + { + name: "Fields empty", + input: NamespaceSummary{VectorCount: 0}, + want: `{"vector_count":0}`, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(c *testing.T) { + got, err := json.Marshal(tt.input) + if err != nil { + c.Errorf("Failed to marshal NamespaceSummary: %v", err) + return + } + if string(got) != tt.want { + c.Errorf("Marshal NamespaceSummary got = %s, want = %s", string(got), tt.want) + } + }) + } +} + +func TestMarshalUsage(t *testing.T) { + tests := []struct { + name string + input Usage + want string + }{ + { + name: "All fields present", + input: Usage{ReadUnits: 100}, + want: `{"read_units":100}`, + }, + { + name: "Fields omitted", + input: Usage{}, + want: `{"read_units":0}`, + }, + { + name: "Fields empty", + input: Usage{ReadUnits: 0}, + want: `{"read_units":0}`, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(c *testing.T) { + got, err := json.Marshal(tt.input) + if err != nil { + c.Errorf("Failed to marshal Usage: %v", err) + return + } + if string(got) != tt.want { + c.Errorf("Marshal Usage got = %s, want = %s", string(got), tt.want) + } + }) + } + +}