leaf/index_test.go

package leaf

import (
	"context"
	"testing"

	"github.com/stretchr/testify/assert"
	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/mongo"
	"go.mongodb.org/mongo-driver/mongo/integration/mtest"
	"go.mongodb.org/mongo-driver/mongo/options"
)

func TestEnsureIndexes(t *testing.T) {
	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
	ctx := context.Background()

	mt.Run("error getting index names", func(mt *mtest.T) {
		mt.AddMockResponses(bson.D{{Key: "ok", Value: 0}})
		_, err := ensureIndexes(ctx, mt.Coll, nil, false)
		assert.Error(t, err)
	})

	mt.Run("no indexes to create", func(mt *mtest.T) {
		mt.AddMockResponses(mtest.CreateCursorResponse(0, "test.foo", mtest.FirstBatch))
		messages, err := ensureIndexes(ctx, mt.Coll, nil, false)
		assert.NoError(t, err)
		assert.Empty(t, messages)
	})

	mt.Run("createIndexes is false", func(mt *mtest.T) {
		indexModel := mongo.IndexModel{
			Keys:    bson.D{bson.E{Key: "test", Value: 1}},
			Options: options.Index().SetName("test_1"),
		}

		mt.AddMockResponses(mtest.CreateCursorResponse(0, "test.foo", mtest.FirstBatch))
		messages, err := ensureIndexes(ctx, mt.Coll, []mongo.IndexModel{indexModel}, false)
		assert.NoError(t, err)
		assert.Equal(t, 1, len(messages))
		assert.Equal(t, "test_1", messages[0].Name)
		assert.Equal(t, "index declared but not present", messages[0].Message)
	})

	mt.Run("createIndexes is true", func(mt *mtest.T) {
		indexModel := mongo.IndexModel{
			Keys:    bson.D{bson.E{Key: "test", Value: 1}},
			Options: options.Index().SetName("test_1"),
		}

		mt.AddMockResponses(mtest.CreateCursorResponse(0, "test.foo", mtest.FirstBatch))
		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.E{Key: "ok", Value: 1}))
		messages, err := ensureIndexes(ctx, mt.Coll, []mongo.IndexModel{indexModel}, true)
		assert.NoError(t, err)
		assert.Equal(t, 1, len(messages))
		assert.Equal(t, "test_1", messages[0].Name)
		assert.Equal(t, "index created", messages[0].Message)
	})
}

func TestIndexNames(t *testing.T) {
	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
	ctx := context.Background()

	mt.Run("test index names", func(mt *mtest.T) {
		collection := mt.Coll

		indexModel := mongo.IndexModel{
			Keys:    bson.D{bson.E{Key: "test", Value: 1}},
			Options: options.Index().SetName("test_1"),
		}

		indexReposne := bson.D{
			{Key: "v", Value: 2},
			{Key: "key", Value: bson.M{"test": 1}},
			{Key: "name", Value: "test_1"},
		}

		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.E{Key: "ok", Value: 1}))
		_, err := collection.Indexes().CreateOne(ctx, indexModel)
		assert.NoError(t, err)

		mt.AddMockResponses(mtest.CreateCursorResponse(0, "test.foo", mtest.FirstBatch, indexReposne))
		names, err := indexNames(ctx, collection)
		assert.NoError(t, err)
		assert.Contains(t, names, "test_1")
	})
}

func TestFilterIndexes(t *testing.T) {
	models := []mongo.IndexModel{
		{
			Keys:    bson.D{bson.E{Key: "test1", Value: 1}},
			Options: options.Index().SetName("test_1"),
		},
		{
			Keys:    bson.D{bson.E{Key: "test2", Value: 1}},
			Options: options.Index(),
		},
		{
			Keys:    bson.D{bson.E{Key: "test3", Value: 1}},
			Options: options.Index().SetName("test_3"),
		},
	}

	names := map[string]struct{}{"test_1": {}}
	filtered, messages := filterIndexes(models, names, []IndexMessage{})

	assert.Equal(t, 1, len(filtered))
	assert.Equal(t, "test_3", *filtered[0].Options.Name)
	assert.Equal(t, 1, len(messages))
	assert.Equal(t, "must specify index name explicitly", messages[0].Message)
}

func TestWithTimestampIndexes(t *testing.T) {
	models := []mongo.IndexModel{
		{
			Keys:    bson.D{bson.E{Key: "test1", Value: 1}},
			Options: options.Index().SetName("test_1"),
		},
	}

	result := withTimestampIndexes(models)

	assert.Equal(t, 3, len(result))
	assert.Equal(t, "test_1", *result[0].Options.Name)
	assert.Equal(t, "created_at_1", *result[1].Options.Name)
	assert.Equal(t, "updated_at_1", *result[2].Options.Name)
}