From cbb7f756af41729c7ddadfc1b7012e7aec77cca8 Mon Sep 17 00:00:00 2001 From: Austin DeNoble Date: Tue, 7 May 2024 18:31:10 -0400 Subject: [PATCH] Add `json:""` marshaling annotations to response structs (#21) ## Problem There's been some UX feedback regarding marshaling of some of our response structs. Primarily, we're not using `json:"key_name"` annotation for our struct fields so they default to capitalized keys in JSON output when marshaled. IE: `{ "Usage": { "ReadUnits": 5 }}`. We'd like these lower case and matched with our API spec. There are also instances where we're returning pointers for certain `*int32` fields which can be irritating to consumers needing to dereference and nil-check. In certain cases we can default any nil pointer values to 0 before returning things. There's no testing for the marshaling on any of our custom structs. ## Solution I may have gone a bit overboard with the testing here. The bulk of this is tests and the tests are all very similar, definitely open to opinions on removing if people find them unnecessary, but changing up how the JSON is produced is pretty easy and it's an important point of integration in my mind, so maybe worth testing for. I also wasn't clear if adding these annotations to our request structs made sense either, so I left those out for now. Overall I'm still working on getting a better understanding of Go syntax and expectations. - Add JSON marshaling annotations to the response structs in `models.go` and `index_connection.go`. - Update `Collection` and `Usage` structs to default integer values rather than using pointers. - Add new `models_test.go` file to run through some validation on struct marshaling, also add unit tests to `index_connection_test.go`. Originally I thought `omitempty` for some of these pointers may be helpful from a consumer standpoint, but I'm actually not sure the more I've thought about it. Does it make sense to omit more or less things? Other feedback here would be much appreciated. ## Type of Change - [ ] Bug fix (non-breaking change which fixes an issue) - [ ] New feature (non-breaking change which adds functionality) - [X] Breaking change (fix or feature that would cause existing functionality to not work as expected) - [ ] This change requires a documentation update - [ ] Infrastructure change (CI configs, etc) - [ ] Non-code change (docs, etc) - [ ] None of the above: (explain here) ## Test Plan Unit tests to explicitly stress marshaling the structs via `json.Marshal()`. --- pinecone/client.go | 14 +- pinecone/index_connection.go | 45 +-- pinecone/index_connection_test.go | 214 +++++++++++- pinecone/models.go | 74 ++-- pinecone/models_test.go | 557 ++++++++++++++++++++++++++++++ 5 files changed, 837 insertions(+), 67 deletions(-) create mode 100644 pinecone/models_test.go 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) + } + }) + } + +}