unit test for bleve prefix query matching

standalone_test.go

package main

import (
	"github.com/blevesearch/bleve"
	"github.com/blevesearch/bleve/analysis/lang/en"
	"strings"
	"testing"
)

type Book struct {
	ID    string
	Title string
}

func (b *Book) Type() string {
	return "book"
}

func TestPrefix(t *testing.T) {
	books := []*Book{
		{Title: "Animal Farm", ID: "1"},
		{Title: "A Moveable Feast", ID: "2"},
		{Title: "Abstract Art", ID: "3"},
		{Title: "The Sun Also Rises", ID: "4"},
		{Title: "Before the Dawn", ID: "5"},
		{Title: "Before the Day is Done", ID: "6"},
		{Title: "Begin Again", ID: "7"},
	}

	idx := setupIndexing(books, t)

	var tests = []struct {
		prefix      string
		expectedIds []string
	}{
		{"a", []string{"1", "2", "3"}},
		{"abstract", []string{"3"}},
		{"be", []string{"5", "6", "7"}},
		{"before", []string{"5", "6"}},
		{"begin", []string{"7"}},
		{"begin aga", []string{"7"}},
		{"bogus", []string{}},
	}

	for _, test := range tests {
		sr, err := PrefixSearch(test.prefix, "Title", idx)
		if err != nil {
			t.Fatal(err)
		}

		if len(sr.Hits) != len(test.expectedIds) {
			t.Errorf("search: '%s' expected %d hits, got %d", test.prefix, len(test.expectedIds), len(sr.Hits))
		}

		for _, id := range test.expectedIds {
			if !idInHits(id, sr) {
				t.Errorf("search: '%s' failed to match id: %s", test.prefix, id)
			}
		}
	}
}

// returns true if the given id appears in the search results
func idInHits(id string, sr *bleve.SearchResult) bool {
	for _, hit := range sr.Hits {
		if hit.ID == id {
			return true
		}
	}

	return false
}

func PrefixSearch(prefix, field string, index bleve.Index) (*bleve.SearchResult, error) {
	prefix = strings.ToLower(prefix)

	prefixQuery := bleve.NewPrefixQuery(prefix)
	prefixQuery.SetField(field)
	req := bleve.NewSearchRequest(prefixQuery)

	// specify which stored fields to return
	req.Fields = []string{field}

	return index.Search(req)
}

func setupIndexing(books []*Book, t *testing.T) bleve.Index {
	t.Helper()

	// a stored generic reusable mapping for english text
	storedEnglishTextFieldMapping := bleve.NewTextFieldMapping()
	storedEnglishTextFieldMapping.Store = true
	storedEnglishTextFieldMapping.Analyzer = en.AnalyzerName

	// only index the fields we specify
	bookMapping := bleve.NewDocumentStaticMapping()
	disabledSection := bleve.NewDocumentDisabledMapping()

	// reduce overall space usage of the index
	// https://groups.google.com/d/msg/bleve/FIJi8Thg-p0/_5QXMhO_CgAJ
	bookMapping.AddSubDocumentMapping("_all", disabledSection)

	// books mapping
	//
	bookMapping.AddFieldMappingsAt("Title", storedEnglishTextFieldMapping)

	// overall mapping
	//
	mapping := bleve.NewIndexMapping()
	mapping.AddDocumentMapping("book", bookMapping)

	idx, err := bleve.NewMemOnly(mapping)
	if err != nil {
		t.Fatal(err)
	}

	for _, book := range books {
		idx.Index(book.ID, book)

	}

	return idx
}