From aebbc204c5563c41ea58a78732db1977f265ec22 Mon Sep 17 00:00:00 2001 From: Preston Vasquez Date: Tue, 17 Sep 2024 11:34:32 -0600 Subject: [PATCH 1/6] Containerize mongovector tests --- vectorstores/mongovector/mongovector_test.go | 178 +++++++++++-------- 1 file changed, 106 insertions(+), 72 deletions(-) diff --git a/vectorstores/mongovector/mongovector_test.go b/vectorstores/mongovector/mongovector_test.go index 8084aec18..9f0960d5a 100644 --- a/vectorstores/mongovector/mongovector_test.go +++ b/vectorstores/mongovector/mongovector_test.go @@ -2,7 +2,6 @@ package mongovector import ( "context" - "errors" "flag" "fmt" "os" @@ -12,6 +11,8 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "github.com/testcontainers/testcontainers-go" + "github.com/testcontainers/testcontainers-go/wait" "github.com/tmc/langchaingo/embeddings" "github.com/tmc/langchaingo/schema" "github.com/tmc/langchaingo/vectorstores" @@ -26,6 +27,7 @@ import ( var testWithoutSetup = flag.Bool("no-atlas-setup", false, "don't create required indexes") const ( + testURI = "MONGODB_VECTOR_TEST_URI" testDB = "langchaingo-test" testColl = "vstore" testIndexDP1536 = "vector_index_dotProduct_1536" @@ -35,34 +37,112 @@ const ( testIndexSize3 = 3 ) -func TestMain(m *testing.M) { - flag.Parse() - - defer func() { - os.Exit(m.Run()) - }() +type atlasContainer struct { + testcontainers.Container + URI string +} - if *testWithoutSetup { - return +func setupAtlas(ctx context.Context) (*atlasContainer, error) { + req := testcontainers.ContainerRequest{ + Image: "mongodb/mongodb-atlas-local", + ExposedPorts: []string{"27017/tcp"}, + WaitingFor: wait.ForLog("Waiting for connections").WithStartupTimeout(1 * time.Second), } - // Create the required vector search indexes for the tests. + container, err := testcontainers.GenericContainer(ctx, testcontainers.GenericContainerRequest{ + ContainerRequest: req, + Started: true, + }) - ctx, cancel := context.WithTimeout(context.Background(), 5*time.Minute) - defer cancel() + if err != nil { + return nil, err + } - if err := resetForE2E(ctx, testIndexDP1536, testIndexSize1536, nil); err != nil { - fmt.Fprintf(os.Stderr, "setup failed for 1536: %v\n", err) + var atlasC *atlasContainer + if container != nil { + atlasC = &atlasContainer{Container: container} } - filters := []string{"pageContent"} - if err := resetForE2E(ctx, testIndexDP1536WithFilter, testIndexSize1536, filters); err != nil { - fmt.Fprintf(os.Stderr, "setup failed for 1536 w filter: %v\n", err) + ip, err := container.Host(ctx) + if err != nil { + return atlasC, err } - if err := resetForE2E(ctx, testIndexDP3, testIndexSize3, nil); err != nil { - fmt.Fprintf(os.Stderr, "setup failed for 3: %v\n", err) + mappedPort, err := container.MappedPort(ctx, "27017") + if err != nil { + return atlasC, err } + + atlasC.URI = fmt.Sprintf("mongodb://%s:%s/?directConnection=true", ip, mappedPort.Port()) + + return atlasC, nil +} + +// resetVectorStore will reset the vector space defined by the given collection. +func resetVectorStore(t *testing.T, coll *mongo.Collection) { + t.Helper() + + filter := bson.D{{Key: pageContentName, Value: bson.D{{Key: "$exists", Value: true}}}} + + _, err := coll.DeleteMany(context.Background(), filter) + assert.NoError(t, err, "failed to reset vector store") +} + +// setupTest will prepare the Atlas vector search for adding to and searching +// a vector space. +func setupTest(t *testing.T, dim int, index string) Store { + t.Helper() + + uri := os.Getenv(testURI) + if uri == "" { + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Minute) + defer cancel() + + container, err := setupAtlas(ctx) + require.NoError(t, err) + + uri = container.URI + os.Setenv(testURI, uri) + } + + require.NotEmpty(t, uri, "URI required") + + client, err := mongo.Connect(options.Client().ApplyURI(uri)) + require.NoError(t, err, "failed to connect to MongoDB server") + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + err = client.Ping(ctx, nil) + + require.NoError(t, err, "failed to ping server") + + time.Sleep(5 * time.Second) // Let the container warm up + + ctx, cancel = context.WithTimeout(context.Background(), 5*time.Minute) + defer cancel() + + err = resetForE2E(ctx, client, testIndexDP1536, testIndexSize1536, nil) + require.NoError(t, err) + + filters := []string{"pageContent"} + err = resetForE2E(ctx, client, testIndexDP1536WithFilter, testIndexSize1536, filters) + require.NoError(t, err) + + err = resetForE2E(ctx, client, testIndexDP3, testIndexSize3, nil) + require.NoError(t, err) + + // Create the vectorstore collection + err = client.Database(testDB).CreateCollection(context.Background(), testColl) + require.NoError(t, err, "failed to create collection") + + coll := client.Database(testDB).Collection(testColl) + resetVectorStore(t, coll) + + emb := newMockEmbedder(dim) + store := New(coll, emb, WithIndex(index)) + + return store } func TestNew(t *testing.T) { @@ -122,44 +202,6 @@ func TestNew(t *testing.T) { } } -// resetVectorStore will reset the vector space defined by the given collection. -func resetVectorStore(t *testing.T, coll *mongo.Collection) { - t.Helper() - - filter := bson.D{{Key: pageContentName, Value: bson.D{{Key: "$exists", Value: true}}}} - - _, err := coll.DeleteMany(context.Background(), filter) - assert.NoError(t, err, "failed to reset vector store") -} - -// setupTest will prepare the Atlas vector search for adding to and searching -// a vector space. -func setupTest(t *testing.T, dim int, index string) Store { - t.Helper() - - uri := os.Getenv("MONGODB_URI") - if uri == "" { - t.Skip("Must set MONGODB_URI to run test") - } - - require.NotEmpty(t, uri, "MONGODB_URI required") - - client, err := mongo.Connect(options.Client().ApplyURI(uri)) - require.NoError(t, err, "failed to connect to MongoDB server") - - // Create the vectorstore collection - err = client.Database(testDB).CreateCollection(context.Background(), testColl) - require.NoError(t, err, "failed to create collection") - - coll := client.Database(testDB).Collection(testColl) - resetVectorStore(t, coll) - - emb := newMockEmbedder(dim) - store := New(coll, emb, WithIndex(index)) - - return store -} - //nolint:paralleltest func TestStore_AddDocuments(t *testing.T) { store := setupTest(t, testIndexSize1536, testIndexDP1536) @@ -536,27 +578,19 @@ func searchIndexExists(ctx context.Context, coll *mongo.Collection, idx string) return false, fmt.Errorf("failed to list search indexes: %w", err) } + if cursor == nil || cursor.Current == nil { + return false, nil + } + name := cursor.Current.Lookup("name").StringValue() queryable := cursor.Current.Lookup("queryable").Boolean() return name == idx && queryable, nil } -func resetForE2E(ctx context.Context, idx string, dim int, filters []string) error { - uri := os.Getenv("MONGODB_URI") - if uri == "" { - return errors.New("MONGODB_URI required") - } - - client, err := mongo.Connect(options.Client().ApplyURI(uri)) - if err != nil { - return fmt.Errorf("failed to connect to server: %w", err) - } - - defer func() { _ = client.Disconnect(ctx) }() - +func resetForE2E(ctx context.Context, client *mongo.Client, idx string, dim int, filters []string) error { // Create the vectorstore collection - err = client.Database(testDB).CreateCollection(ctx, testColl) + err := client.Database(testDB).CreateCollection(ctx, testColl) if err != nil { return fmt.Errorf("failed to create vector store collection: %w", err) } @@ -585,7 +619,7 @@ func resetForE2E(ctx context.Context, idx string, dim int, filters []string) err _, err = createVectorSearchIndex(ctx, coll, idx, fields...) if err != nil { - return fmt.Errorf("faield to create index: %w", err) + return fmt.Errorf("failed to create index: %w", err) } return nil From 80f05b8dbb87bbb73d756d9f0ffab0ea782a5cdb Mon Sep 17 00:00:00 2001 From: Preston Vasquez Date: Tue, 17 Sep 2024 11:44:09 -0600 Subject: [PATCH 2/6] Resolve linting errors --- vectorstores/mongovector/mongovector_test.go | 19 ++++++++++++------- 1 file changed, 12 insertions(+), 7 deletions(-) diff --git a/vectorstores/mongovector/mongovector_test.go b/vectorstores/mongovector/mongovector_test.go index 9f0960d5a..67aff2733 100644 --- a/vectorstores/mongovector/mongovector_test.go +++ b/vectorstores/mongovector/mongovector_test.go @@ -2,8 +2,9 @@ package mongovector import ( "context" - "flag" "fmt" + "net" + "net/url" "os" "strings" "testing" @@ -21,11 +22,6 @@ import ( "go.mongodb.org/mongo-driver/v2/mongo/options" ) -// Run the test without setting up the test space. -// -//nolint:gochecknoglobals -var testWithoutSetup = flag.Bool("no-atlas-setup", false, "don't create required indexes") - const ( testURI = "MONGODB_VECTOR_TEST_URI" testDB = "langchaingo-test" @@ -73,7 +69,16 @@ func setupAtlas(ctx context.Context) (*atlasContainer, error) { return atlasC, err } - atlasC.URI = fmt.Sprintf("mongodb://%s:%s/?directConnection=true", ip, mappedPort.Port()) + uri := &url.URL{ + Scheme: "mongodb", + Host: net.JoinHostPort(ip, mappedPort.Port()), + Path: "/", + RawQuery: "directConnection=true", + } + + fmt.Println(uri.String()) + + atlasC.URI = uri.String() return atlasC, nil } From f46f0dea1c08670055a45cd172e6e3aefa4a9304 Mon Sep 17 00:00:00 2001 From: Preston Vasquez Date: Tue, 17 Sep 2024 11:47:05 -0600 Subject: [PATCH 3/6] gofumpt test file --- vectorstores/mongovector/mongovector_test.go | 1 - 1 file changed, 1 deletion(-) diff --git a/vectorstores/mongovector/mongovector_test.go b/vectorstores/mongovector/mongovector_test.go index 67aff2733..a474b0240 100644 --- a/vectorstores/mongovector/mongovector_test.go +++ b/vectorstores/mongovector/mongovector_test.go @@ -49,7 +49,6 @@ func setupAtlas(ctx context.Context) (*atlasContainer, error) { ContainerRequest: req, Started: true, }) - if err != nil { return nil, err } From 48fc115df816ff7dc4cdb87821fac6f966d1ebd1 Mon Sep 17 00:00:00 2001 From: Preston Vasquez Date: Tue, 17 Sep 2024 11:56:14 -0600 Subject: [PATCH 4/6] Remove empty line --- vectorstores/mongovector/mongovector_test.go | 1 - 1 file changed, 1 deletion(-) diff --git a/vectorstores/mongovector/mongovector_test.go b/vectorstores/mongovector/mongovector_test.go index a474b0240..e581c6646 100644 --- a/vectorstores/mongovector/mongovector_test.go +++ b/vectorstores/mongovector/mongovector_test.go @@ -118,7 +118,6 @@ func setupTest(t *testing.T, dim int, index string) Store { defer cancel() err = client.Ping(ctx, nil) - require.NoError(t, err, "failed to ping server") time.Sleep(5 * time.Second) // Let the container warm up From f8140daf5334b54f64c7bc5e52b43135b3a7bbf6 Mon Sep 17 00:00:00 2001 From: Preston Vasquez Date: Wed, 18 Sep 2024 09:06:02 -0600 Subject: [PATCH 5/6] GODRIVER-3345 Remove print statement --- vectorstores/mongovector/mongovector_test.go | 2 -- 1 file changed, 2 deletions(-) diff --git a/vectorstores/mongovector/mongovector_test.go b/vectorstores/mongovector/mongovector_test.go index e581c6646..f97b50616 100644 --- a/vectorstores/mongovector/mongovector_test.go +++ b/vectorstores/mongovector/mongovector_test.go @@ -75,8 +75,6 @@ func setupAtlas(ctx context.Context) (*atlasContainer, error) { RawQuery: "directConnection=true", } - fmt.Println(uri.String()) - atlasC.URI = uri.String() return atlasC, nil From eff62dc27aa0ef922fbaea53f42334ff34dd657e Mon Sep 17 00:00:00 2001 From: Preston Vasquez Date: Thu, 19 Sep 2024 15:00:05 -0600 Subject: [PATCH 6/6] Increase warm up time --- vectorstores/mongovector/mongovector_test.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vectorstores/mongovector/mongovector_test.go b/vectorstores/mongovector/mongovector_test.go index f97b50616..6cb17b5a3 100644 --- a/vectorstores/mongovector/mongovector_test.go +++ b/vectorstores/mongovector/mongovector_test.go @@ -118,7 +118,7 @@ func setupTest(t *testing.T, dim int, index string) Store { err = client.Ping(ctx, nil) require.NoError(t, err, "failed to ping server") - time.Sleep(5 * time.Second) // Let the container warm up + time.Sleep(10 * time.Second) // Let the container warm up ctx, cancel = context.WithTimeout(context.Background(), 5*time.Minute) defer cancel()