Add Pinecone Local testing support in CI (#77)

## Problem
We want to add CI that validates testing the Go SDK against an instance
of pclocal hosted in a Docker container.

## Solution
- Update the `ci.yaml` GitHub workflow file to add a new `services`
block, which includes using the
`ghcr.io/pinecone-io/pinecone-index:latest` Docker image to create and
how two indexes locally (pods and serverless).
- Add new `/pinecone/local_test.go` which defines a new
`LocalIntegrationTests` struct, which isolates a collection of tests
specifically meant for pclocal. The tests are isolated from our unit +
integration test suite with the `//go:build localServer` build tag. Two
suites are defined for pods and serverless, and each set of tests is run
against both indexes. For now the test suite generates 100 vectors,
upserts them to the index, tests basic Fetch, Query, Update,
DescribeIndexStats, and Delete.
- Add a new `Run local integration tests` step to the `ci.yaml`
workflow. This triggers the new `local-test.go` file which tests both
instances of locally hosted indexes (`PINECONE_INDEX_URL_POD`,
`PINECONE_INDEX_URL_SERVERLESS`) with `go test -count=1 -v ./pinecone
-run TestRunLocalIntegrationSuite -tags=localServer`.

### To Do
I'd like to expand the number of different tests we exercise on each
index. Specifically:
- Metadata filtering
- Sparse vectors
- Multiple namespaces
- List vectors

## Type of Change
- [ ] Bug fix (non-breaking change which fixes an issue)
- [ ] New feature (non-breaking change which adds functionality)
- [ ] Breaking change (fix or feature that would cause existing
functionality to not work as expected)
- [ ] This change requires a documentation update
- [X] Infrastructure change (CI configs, etc)
- [ ] Non-code change (docs, etc)
- [ ] None of the above: (explain here)

## Test Plan
Make sure the CI workflow passes for this PR.

If you'd like to run the local integration tests locally, you will need
to download Docker, and use `docker compose up` to create the same
indexes on your local machine. Once you've done that, you'll need to
export environment variables for the different index addresses, along
with the dimension.

Here's an example:

```bash
export PINECONE_INDEX_URL_POD="http://localhost:5082" PINECONE_INDEX_URL_SERVERLESS="http://localhost:5081" PINECONE_DIMENSION="1536"
```

Then, you can run the local integration tests directly:
```bash
go test -count=1 -v ./pinecone -run TestRunLocalIntegrationSuite -tags=localServer
```

---
- To see the specific tasks where the Asana app for GitHub is being
used, see below:
  - https://app.asana.com/0/0/1208056331944375

.github/workflows/ci.yaml

@@ -7,6 +7,25 @@ on:
 jobs:
  build:
  runs-on: ubuntu-latest
+ services:
+ pc-index-serverless:
+ image: ghcr.io/pinecone-io/pinecone-index:latest
+ ports:
+ - 5081:5081
+ env:
+ PORT: 5081
+ DIMENSION: 1536
+ METRIC: dot-product
+ INDEX_TYPE: serverless
+ pc-index-pod:
+ image: ghcr.io/pinecone-io/pinecone-index:latest
+ ports:
+ - 5082:5082
+ env:
+ PORT: 5082
+ DIMENSION: 1536
+ METRIC: cosine
+ INDEX_TYPE: pod
  steps:
  - uses: actions/checkout@v4
  - name: Setup Go
@@ -17,6 +36,13 @@ jobs:
  run: |
  go get ./pinecone
  - name: Run tests
+ continue-on-error: true
  run: go test -count=1 -v ./pinecone
  env:
  PINECONE_API_KEY: ${{ secrets.API_KEY }}
+ - name: Run local integration tests
+ run: go test -count=1 -v ./pinecone -run TestRunLocalIntegrationSuite -tags=localServer
+ env:
+ PINECONE_INDEX_URL_POD: http://localhost:5082
+ PINECONE_INDEX_URL_SERVERLESS: http://localhost:5081
+ PINECONE_DIMENSION: 1536

pinecone/index_connection.go

@@ -389,7 +389,7 @@ func (idx *IndexConnection) ListVectors(ctx context.Context, in *ListVectorsRequ
 
  return &ListVectorsResponse{
  VectorIds: vectorIds,
- Usage: &Usage{ReadUnits: derefOrDefault(res.Usage.ReadUnits, 0)},
+ Usage: toUsage(res.Usage),
  NextPaginationToken: toPaginationToken(res.Pagination),
  Namespace: idx.Namespace,
  }, nil

pinecone/index_connection_test.go

@@ -61,7 +61,7 @@ func (ts *IntegrationTests) TestDeleteVectorsById() {
  assert.NoError(ts.T(), err)
  ts.vectorIds = []string{}
 
- vectors := GenerateVectors(5, ts.dimension, true)
+ vectors := GenerateVectors(5, ts.dimension, true, nil)
 
  _, err = ts.idxConn.UpsertVectors(ctx, vectors)
  if err != nil {
@@ -96,7 +96,7 @@ func (ts *IntegrationTests) TestDeleteVectorsByFilter() {
  }
  ts.vectorIds = []string{}
 
- vectors := GenerateVectors(5, ts.dimension, true)
+ vectors := GenerateVectors(5, ts.dimension, true, nil)
 
  _, err = ts.idxConn.UpsertVectors(ctx, vectors)
  if err != nil {
@@ -117,7 +117,7 @@ func (ts *IntegrationTests) TestDeleteAllVectorsInNamespace() {
  assert.NoError(ts.T(), err)
  ts.vectorIds = []string{}
 
- vectors := GenerateVectors(5, ts.dimension, true)
+ vectors := GenerateVectors(5, ts.dimension, true, nil)
 
  _, err = ts.idxConn.UpsertVectors(ctx, vectors)
  if err != nil {

pinecone/local_test.go

@@ -0,0 +1,246 @@
+//go:build localServer
+
+package pinecone
+
+import (
+ "context"
+ "fmt"
+ "os"
+ "strconv"
+ "testing"
+
+ "github.com/stretchr/testify/assert"
+ "github.com/stretchr/testify/require"
+ "github.com/stretchr/testify/suite"
+ "google.golang.org/grpc"
+ "google.golang.org/grpc/credentials/insecure"
+ "google.golang.org/protobuf/encoding/protojson"
+ "google.golang.org/protobuf/types/known/structpb"
+)
+
+type LocalIntegrationTests struct {
+ suite.Suite
+ client *Client
+ host string
+ dimension int32
+ indexType string
+ namespace string
+ metadata *Metadata
+ vectorIds []string
+ idxConns []*IndexConnection
+}
+
+func (ts *LocalIntegrationTests) SetupSuite() {
+ ctx := context.Background()
+
+ // Deterministically create vectors
+ vectors := GenerateVectors(100, ts.dimension, true, ts.metadata)
+
+ // Get vector ids for the suite
+ vectorIds := make([]string, len(vectors))
+ for i, v := range vectors {
+ vectorIds[i] = v.Id
+ }
+
+ // Upsert vectors into each index connection
+ for _, idxConn := range ts.idxConns {
+ upsertedVectors, err := idxConn.UpsertVectors(ctx, vectors)
+ require.NoError(ts.T(), err)
+ fmt.Printf("Upserted vectors: %v into host: %s in namespace: %s \n", upsertedVectors, ts.host, idxConn.Namespace)
+ }
+
+ ts.vectorIds = append(ts.vectorIds, vectorIds...)
+}
+
+func (ts *LocalIntegrationTests) TearDownSuite() {
+ // test deleting vectors as a part of cleanup for each index connection
+ for _, idxConn := range ts.idxConns {
+ // Delete a slice of vectors by id
+ err := idxConn.DeleteVectorsById(context.Background(), ts.vectorIds[10:20])
+ require.NoError(ts.T(), err)
+
+ // Delete vectors by filter
+ if ts.indexType == "pods" {
+ err = idxConn.DeleteVectorsByFilter(context.Background(), ts.metadata)
+ require.NoError(ts.T(), err)
+ }
+
+ // Delete all remaining vectors
+ err = idxConn.DeleteAllVectorsInNamespace(context.Background())
+ require.NoError(ts.T(), err)
+ }
+
+ description, err := ts.idxConns[0].DescribeIndexStats(context.Background())
+ require.NoError(ts.T(), err)
+ assert.NotNil(ts.T(), description, "Index description should not be nil")
+ assert.Equal(ts.T(), uint32(0), description.TotalVectorCount, "Total vector count should be 0 after deleting")
+}
+
+// This is the entry point for all local integration tests
+// This test function is picked up by go test and triggers the suite runs when
+// the build tag localServer is set
+func TestRunLocalIntegrationSuite(t *testing.T) {
+ fmt.Println("Running local integration tests")
+ RunLocalSuite(t)
+}
+
+func RunLocalSuite(t *testing.T) {
+ fmt.Println("Running local integration tests")
+ localHostPod, present := os.LookupEnv("PINECONE_INDEX_URL_POD")
+ assert.True(t, present, "PINECONE_INDEX_URL_POD env variable not set")
+
+ localHostServerless, present := os.LookupEnv("PINECONE_INDEX_URL_SERVERLESS")
+ assert.True(t, present, "PINECONE_INDEX_URL_SERVERLESS env variable not set")
+
+ dimension, present := os.LookupEnv("PINECONE_DIMENSION")
+ assert.True(t, present, "PINECONE_DIMENSION env variable not set")
+
+ parsedDimension, err := strconv.ParseInt(dimension, 10, 32)
+ require.NoError(t, err)
+
+ namespace := "test-namespace"
+ metadata := &structpb.Struct{
+ Fields: map[string]*structpb.Value{
+ "genre": {Kind: &structpb.Value_StringValue{StringValue: "classical"}},
+ },
+ }
+
+ client, err := NewClientBase(NewClientBaseParams{})
+ require.NotNil(t, client, "Client should not be nil after creation")
+ require.NoError(t, err)
+
+ // Create index connections for pod and serverless indexes with both default namespace
+ // and a custom namespace
+ var podIdxConns []*IndexConnection
+ idxConnPod, err := client.Index(NewIndexConnParams{Host: localHostPod})
+ require.NoError(t, err)
+ podIdxConns = append(podIdxConns, idxConnPod)
+
+ idxConnPodNamespace, err := client.Index(NewIndexConnParams{Host: localHostPod, Namespace: namespace})
+ require.NoError(t, err)
+ podIdxConns = append(podIdxConns, idxConnPodNamespace)
+
+ var serverlessIdxConns []*IndexConnection
+ idxConnServerless, err := client.Index(NewIndexConnParams{Host: localHostServerless},
+ grpc.WithTransportCredentials(insecure.NewCredentials()))
+ require.NoError(t, err)
+ serverlessIdxConns = append(serverlessIdxConns, idxConnServerless)
+
+ idxConnServerless, err = client.Index(NewIndexConnParams{Host: localHostServerless, Namespace: namespace})
+ require.NoError(t, err)
+ serverlessIdxConns = append(serverlessIdxConns, idxConnServerless)
+
+ localHostPodSuite := &LocalIntegrationTests{
+ client: client,
+ idxConns: podIdxConns,
+ indexType: "pods",
+ host: localHostPod,
+ namespace: namespace,
+ metadata: metadata,
+ dimension: int32(parsedDimension),
+ }
+
+ localHostSuiteServerless := &LocalIntegrationTests{
+ client: client,
+ idxConns: serverlessIdxConns,
+ indexType: "serverless",
+ host: localHostServerless,
+ namespace: namespace,
+ metadata: metadata,
+ dimension: int32(parsedDimension),
+ }
+
+ suite.Run(t, localHostPodSuite)
+ suite.Run(t, localHostSuiteServerless)
+}
+
+func (ts *LocalIntegrationTests) TestFetchVectors() {
+ fetchVectorId := ts.vectorIds[0]
+
+ for _, idxConn := range ts.idxConns {
+ fetchVectorsResponse, err := idxConn.FetchVectors(context.Background(), []string{fetchVectorId})
+ require.NoError(ts.T(), err)
+
+ assert.NotNil(ts.T(), fetchVectorsResponse, "Fetch vectors response should not be nil")
+ assert.Equal(ts.T(), 1, len(fetchVectorsResponse.Vectors), "Fetch vectors response should have 1 vector")
+ assert.Equal(ts.T(), fetchVectorId, fetchVectorsResponse.Vectors[fetchVectorId].Id, "Fetched vector id should match")
+ }
+}
+
+func (ts *LocalIntegrationTests) TestQueryVectors() {
+ queryVectorId := ts.vectorIds[0]
+ topK := 10
+
+ for _, idxConn := range ts.idxConns {
+ queryVectorsByIdResponse, err := idxConn.QueryByVectorId(context.Background(), &QueryByVectorIdRequest{
+ VectorId: queryVectorId,
+ TopK: uint32(topK),
+ IncludeValues: true,
+ IncludeMetadata: true,
+ })
+ require.NoError(ts.T(), err)
+
+ assert.NotNil(ts.T(), queryVectorsByIdResponse, "QueryByVectorId results should not be nil")
+ assert.Equal(ts.T(), topK, len(queryVectorsByIdResponse.Matches), "QueryByVectorId results should have 10 matches")
+ assert.Equal(ts.T(), queryVectorId, queryVectorsByIdResponse.Matches[0].Vector.Id, "Top QueryByVectorId result's vector id should match queryVectorId")
+
+ queryByVectorValuesResponse, err := idxConn.QueryByVectorValues(context.Background(), &QueryByVectorValuesRequest{
+ Vector: queryVectorsByIdResponse.Matches[0].Vector.Values,
+ TopK: uint32(topK),
+ MetadataFilter: ts.metadata,
+ IncludeValues: true,
+ IncludeMetadata: true,
+ })
+ require.NoError(ts.T(), err)
+
+ assert.NotNil(ts.T(), queryByVectorValuesResponse, "QueryByVectorValues results should not be nil")
+ assert.Equal(ts.T(), topK, len(queryByVectorValuesResponse.Matches), "QueryByVectorValues results should have 10 matches")
+
+ resultMetadata, err := protojson.Marshal(queryByVectorValuesResponse.Matches[0].Vector.Metadata)
+ assert.NoError(ts.T(), err)
+ suiteMetadata, err := protojson.Marshal(ts.metadata)
+ assert.NoError(ts.T(), err)
+
+ assert.Equal(ts.T(), resultMetadata, suiteMetadata, "Top QueryByVectorValues result's metadata should match the test suite's metadata")
+ }
+}
+
+func (ts *LocalIntegrationTests) TestUpdateVectors() {
+ updateVectorId := ts.vectorIds[0]
+ newValues := generateVectorValues(ts.dimension)
+
+ for _, idxConn := range ts.idxConns {
+ err := idxConn.UpdateVector(context.Background(), &UpdateVectorRequest{Id: updateVectorId, Values: newValues})
+ require.NoError(ts.T(), err)
+
+ fetchVectorsResponse, err := idxConn.FetchVectors(context.Background(), []string{updateVectorId})
+ require.NoError(ts.T(), err)
+ assert.Equal(ts.T(), newValues, fetchVectorsResponse.Vectors[updateVectorId].Values, "Updated vector values should match")
+ }
+}
+
+func (ts *LocalIntegrationTests) TestDescribeIndexStats() {
+ for _, idxConn := range ts.idxConns {
+ description, err := idxConn.DescribeIndexStats(context.Background())
+ require.NoError(ts.T(), err)
+
+ assert.NotNil(ts.T(), description, "Index description should not be nil")
+ assert.Equal(ts.T(), description.TotalVectorCount, uint32(len(ts.vectorIds)*2), "Index host should match")
+ }
+}
+
+func (ts *LocalIntegrationTests) TestListVectorIds() {
+ limit := uint32(25)
+ // Listing vector ids is only available for serverless indexes
+ if ts.indexType == "serverless" {
+ for _, idxConn := range ts.idxConns {
+ listVectorIdsResponse, err := idxConn.ListVectors(context.Background(), &ListVectorsRequest{
+ Limit: &limit,
+ })
+ require.NoError(ts.T(), err)
+
+ assert.NotNil(ts.T(), listVectorIdsResponse, "ListVectors response should not be nil")
+ assert.Equal(ts.T(), limit, uint32(len(listVectorIdsResponse.VectorIds)), "ListVectors response should have %d vector ids", limit)
+ }
+ }
+}

pinecone/test_suite.go

@@ -7,8 +7,6 @@ import (
  "math/rand"
  "time"
 
- "google.golang.org/protobuf/types/known/structpb"
-
  "github.com/google/uuid"
  "github.com/stretchr/testify/require"
  "github.com/stretchr/testify/suite"
@@ -47,14 +45,13 @@ func (ts *IntegrationTests) SetupSuite() {
  ts.idxConn = idxConn
 
  // Deterministically create vectors
- vectors := GenerateVectors(10, ts.dimension, false)
+ vectors := GenerateVectors(10, ts.dimension, false, nil)
 
  // Add vector ids to the suite
  vectorIds := make([]string, len(vectors))
  for i, v := range vectors {
  vectorIds[i] = v.Id
  }
- ts.vectorIds = append(ts.vectorIds, vectorIds...)
 
  // Upsert vectors
  err = upsertVectors(ts, ctx, vectors)
@@ -158,7 +155,7 @@ func WaitUntilIndexReady(ts *IntegrationTests, ctx context.Context) (bool, error
  }
 }
 
-func GenerateVectors(numOfVectors int, dimension int32, isSparse bool) []*Vector {
+func GenerateVectors(numOfVectors int, dimension int32, isSparse bool, metadata *Metadata) []*Vector {
  vectors := make([]*Vector, numOfVectors)
 
  for i := 0; i < int(numOfVectors); i++ {
@@ -177,12 +174,9 @@ func GenerateVectors(numOfVectors int, dimension int32, isSparse bool) []*Vector
  vectors[i].SparseValues = &sparseValues
  }
 
- metadata := &structpb.Struct{
- Fields: map[string]*structpb.Value{
- "genre": {Kind: &structpb.Value_StringValue{StringValue: "classical"}},
- },
+ if metadata != nil {
+ vectors[i].Metadata = metadata
  }
- vectors[i].Metadata = metadata
  }
 
  return vectors