Skip to content

Commit

Permalink
add list index tests
Browse files Browse the repository at this point in the history
  • Loading branch information
Jesse Schmidt committed Jun 11, 2024
1 parent 0cb992f commit 419a1b5
Show file tree
Hide file tree
Showing 2 changed files with 162 additions and 2 deletions.
6 changes: 4 additions & 2 deletions cmd/listIndex.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -45,8 +45,10 @@ var listIndexRequiredFlags = []string{}

// listIndexCmd represents the listIndex command
func newListIndexCmd() *cobra.Command {
return &cobra.Command{Use: "index",
Short: "A command for listing indexes",
return &cobra.Command{
Use: "index",
Aliases: []string{"indexes"},
Short: "A command for listing indexes",
Long: fmt.Sprintf(`A command for displaying useful information about AVS indexes. To display additional
index information use the --%s flag.
For example:
Expand Down
158 changes: 158 additions & 0 deletions e2e_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ import (
"os"
"os/exec"
"path"
"regexp"
"strings"
"testing"
"time"
Expand Down Expand Up @@ -139,6 +140,7 @@ func getBoolPtr(b bool) *bool {
type IndexDefinitionBuilder struct {
indexName string
namespace string
set *string
dimension int
vectorDistanceMetric protos.VectorDistanceMetric
vectorField string
Expand Down Expand Up @@ -168,6 +170,11 @@ func NewIndexDefinitionBuilder(
}
}

func (idb *IndexDefinitionBuilder) WithSet(set string) *IndexDefinitionBuilder {
idb.set = &set
return idb
}

func (idb *IndexDefinitionBuilder) WithStorageNamespace(storageNamespace string) *IndexDefinitionBuilder {
idb.storageNamespace = &storageNamespace
return idb
Expand Down Expand Up @@ -236,6 +243,10 @@ func (idb *IndexDefinitionBuilder) Build() *protos.IndexDefinition {
},
}

if idb.set != nil {
indexDef.SetFilter = idb.set
}

if idb.storageNamespace != nil {
indexDef.Storage.Namespace = idb.storageNamespace
}
Expand Down Expand Up @@ -390,6 +401,153 @@ func (suite *CmdTestSuite) TestDropIndexFailsDoesNotExistCmd() {
suite.Assert().Contains(lines[0], "server error")
}

func removeANSICodes(input string) string {
re := regexp.MustCompile(`\x1b[^m]*m`)
return re.ReplaceAllString(input, "")
}

func (suite *CmdTestSuite) TestSuccessfulListIndexCmd() {
indexes, err := suite.avsClient.IndexList(context.Background())
if err != nil {
suite.FailNow(err.Error())
}

for _, index := range indexes.GetIndices() {
err := suite.avsClient.IndexDrop(context.Background(), index.Id.Namespace, index.Id.Name)
if err != nil {
suite.FailNow(err.Error())
}
}

testCases := []struct {
name string
indexes []*protos.IndexDefinition
cmd string
expectedTable string
}{
{
"single index",
[]*protos.IndexDefinition{
NewIndexDefinitionBuilder(
"list", "test", 256, protos.VectorDistanceMetric_COSINE, "vector",
).Build(),
},
fmt.Sprintf("list index -h %s", suite.avsHostPort.String()),
`╭─────────────────────────────────────────────────────────────────────────╮
│ Indexes │
├───┬──────┬───────────┬────────┬────────────┬─────────────────┬──────────┤
│ │ NAME │ NAMESPACE │ FIELD │ DIMENSIONS │ DISTANCE METRIC │ UNMERGED │
├───┼──────┼───────────┼────────┼────────────┼─────────────────┼──────────┤
│ 1 │ list │ test │ vector │ 256 │ COSINE │ 0 │
╰───┴──────┴───────────┴────────┴────────────┴─────────────────┴──────────╯
`,
},
{
"double index with set",
[]*protos.IndexDefinition{
NewIndexDefinitionBuilder(
"list1", "test", 256, protos.VectorDistanceMetric_COSINE, "vector",
).Build(),
NewIndexDefinitionBuilder(
"list2", "bar", 256, protos.VectorDistanceMetric_HAMMING, "vector",
).WithSet("barset").Build(),
},
fmt.Sprintf("list index -h %s", suite.avsHostPort.String()),
`╭───────────────────────────────────────────────────────────────────────────────────╮
│ Indexes │
├───┬───────┬───────────┬────────┬────────┬────────────┬─────────────────┬──────────┤
│ │ NAME │ NAMESPACE │ SET │ FIELD │ DIMENSIONS │ DISTANCE METRIC │ UNMERGED │
├───┼───────┼───────────┼────────┼────────┼────────────┼─────────────────┼──────────┤
│ 1 │ list2 │ bar │ barset │ vector │ 256 │ HAMMING │ 0 │
├───┼───────┼───────────┼────────┼────────┼────────────┼─────────────────┼──────────┤
│ 2 │ list1 │ test │ │ vector │ 256 │ COSINE │ 0 │
╰───┴───────┴───────────┴────────┴────────┴────────────┴─────────────────┴──────────╯
`,
},
{
"double index with set and verbose",
[]*protos.IndexDefinition{
NewIndexDefinitionBuilder(
"list1", "test", 256, protos.VectorDistanceMetric_COSINE, "vector",
).Build(),
NewIndexDefinitionBuilder(
"list2", "bar", 256, protos.VectorDistanceMetric_HAMMING, "vector",
).WithSet("barset").Build(),
},
fmt.Sprintf("list index -h %s --verbose", suite.avsHostPort.String()),
`╭────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╮
│ Indexes │
├───┬───────┬───────────┬────────┬────────┬────────────┬─────────────────┬──────────┬───────────────────────┬────────────────────────────────┤
│ │ NAME │ NAMESPACE │ SET │ FIELD │ DIMENSIONS │ DISTANCE METRIC │ UNMERGED │ STORAGE │ INDEX PARAMETERS │
├───┼───────┼───────────┼────────┼────────┼────────────┼─────────────────┼──────────┼───────────────────────┼────────────────────────────────┤
│ 1 │ list2 │ bar │ barset │ vector │ 256 │ HAMMING │ 0 │ ╭───────────┬───────╮ │ ╭────────────────────────────╮ │
│ │ │ │ │ │ │ │ │ │ Namespace │ bar │ │ │ HNSW │ │
│ │ │ │ │ │ │ │ │ │ Set │ list2 │ │ ├───────────────────┬────────┤ │
│ │ │ │ │ │ │ │ │ ╰───────────┴───────╯ │ │ Max Edges │ 16 │ │
│ │ │ │ │ │ │ │ │ │ │ Ef │ 100 │ │
│ │ │ │ │ │ │ │ │ │ │ Construction Ef │ 100 │ │
│ │ │ │ │ │ │ │ │ │ │ Batch Max Records │ 100000 │ │
│ │ │ │ │ │ │ │ │ │ │ Batch Interval │ 30000 │ │
│ │ │ │ │ │ │ │ │ │ │ Batch Enabled │ true │ │
│ │ │ │ │ │ │ │ │ │ ╰───────────────────┴────────╯ │
├───┼───────┼───────────┼────────┼────────┼────────────┼─────────────────┼──────────┼───────────────────────┼────────────────────────────────┤
│ 2 │ list1 │ test │ │ vector │ 256 │ COSINE │ 0 │ ╭───────────┬───────╮ │ ╭────────────────────────────╮ │
│ │ │ │ │ │ │ │ │ │ Namespace │ test │ │ │ HNSW │ │
│ │ │ │ │ │ │ │ │ │ Set │ list1 │ │ ├───────────────────┬────────┤ │
│ │ │ │ │ │ │ │ │ ╰───────────┴───────╯ │ │ Max Edges │ 16 │ │
│ │ │ │ │ │ │ │ │ │ │ Ef │ 100 │ │
│ │ │ │ │ │ │ │ │ │ │ Construction Ef │ 100 │ │
│ │ │ │ │ │ │ │ │ │ │ Batch Max Records │ 100000 │ │
│ │ │ │ │ │ │ │ │ │ │ Batch Interval │ 30000 │ │
│ │ │ │ │ │ │ │ │ │ │ Batch Enabled │ true │ │
│ │ │ │ │ │ │ │ │ │ ╰───────────────────┴────────╯ │
╰───┴───────┴───────────┴────────┴────────┴────────────┴─────────────────┴──────────┴───────────────────────┴────────────────────────────────╯
`,
},
}

for _, tc := range testCases {
suite.Run(tc.name, func() {
for _, index := range tc.indexes {
setFilter := []string{}
if index.SetFilter != nil {
setFilter = append(setFilter, *index.SetFilter)
}

err := suite.avsClient.IndexCreate(
context.Background(),
index.Id.Namespace,
setFilter,
index.Id.Name,
index.GetField(),
index.GetDimensions(),
index.GetVectorDistanceMetric(),
index.GetHnswParams(),
index.GetLabels(),
index.GetStorage(),
)
if err != nil {
suite.FailNowf("unable to create index", "%v", err)
}

defer suite.avsClient.IndexDrop(
context.Background(),
index.Id.Namespace,
index.Id.Name,
)
}

lines, err := suite.runCmd(strings.Split(tc.cmd, " ")...)
suite.Assert().NoError(err, "error: %s, stdout/err: %s", err, lines)

actualTable := removeANSICodes(strings.Join(lines, "\n"))

suite.Assert().Equal(tc.expectedTable, actualTable)

})
}
}

func (suite *CmdTestSuite) TestFailInvalidArg() {
testCases := []struct {
name string
Expand Down

0 comments on commit 419a1b5

Please sign in to comment.