From 699c86073ad1efd3a26d080d7543a14e40ae88ba Mon Sep 17 00:00:00 2001
From: Marty Schoch <marty.schoch@gmail.com>
Date: Tue, 1 Dec 2015 12:29:56 -0500
Subject: [PATCH] make existing integration tests work with firestorm

---
 config.go                    |  2 ++
 index/firestorm/firestorm.go |  5 +++++
 index/firestorm/stats.go     |  2 ++
 index/index.go               | 11 +++++++++++
 test/integration_test.go     | 20 ++++++++++++++++++++
 5 files changed, 40 insertions(+)

diff --git a/config.go b/config.go
index f9100bf2..904ca228 100644
--- a/config.go
+++ b/config.go
@@ -20,6 +20,8 @@ import (
 	"github.com/blevesearch/bleve/index/upside_down"
 	"github.com/blevesearch/bleve/registry"
 	"github.com/blevesearch/bleve/search/highlight/highlighters/html"
+
+	_ "github.com/blevesearch/bleve/index/firestorm"
 )
 
 var bleveExpVar = expvar.NewMap("bleve")
diff --git a/index/firestorm/firestorm.go b/index/firestorm/firestorm.go
index 5dfdec13..839665e2 100644
--- a/index/firestorm/firestorm.go
+++ b/index/firestorm/firestorm.go
@@ -53,6 +53,7 @@ func NewFirestorm(storeName string, storeConfig map[string]interface{}, analysis
 		highDocNumber: 0,
 		stats:         &indexStat{},
 	}
+	rv.stats.f = &rv
 	rv.garbageCollector = NewGarbageCollector(&rv)
 	rv.lookuper = NewLookuper(&rv)
 	rv.dictUpdater = NewDictUpdater(&rv)
@@ -462,6 +463,10 @@ func (f *Firestorm) Stats() json.Marshaler {
 
 }
 
+func (f *Firestorm) Wait(timeout time.Duration) error {
+	return f.dictUpdater.waitTasksDone(timeout)
+}
+
 func init() {
 	registry.RegisterIndexType(Name, NewFirestorm)
 }
diff --git a/index/firestorm/stats.go b/index/firestorm/stats.go
index 2c4fad01..4f3eae87 100644
--- a/index/firestorm/stats.go
+++ b/index/firestorm/stats.go
@@ -15,6 +15,7 @@ import (
 )
 
 type indexStat struct {
+	f                                 *Firestorm
 	updates, deletes, batches, errors uint64
 	analysisTime, indexTime           uint64
 }
@@ -27,5 +28,6 @@ func (i *indexStat) MarshalJSON() ([]byte, error) {
 	m["errors"] = atomic.LoadUint64(&i.errors)
 	m["analysis_time"] = atomic.LoadUint64(&i.analysisTime)
 	m["index_time"] = atomic.LoadUint64(&i.indexTime)
+	m["lookup_queue_len"] = len(i.f.lookuper.workChan)
 	return json.Marshal(m)
 }
diff --git a/index/index.go b/index/index.go
index 04626e4b..4f15b45b 100644
--- a/index/index.go
+++ b/index/index.go
@@ -12,6 +12,7 @@ package index
 import (
 	"encoding/json"
 	"fmt"
+	"time"
 
 	"github.com/blevesearch/bleve/document"
 )
@@ -44,6 +45,16 @@ type Index interface {
 	Analyze(d *document.Document) *AnalysisResult
 }
 
+// AsyncIndex is an interface for indexes which perform
+// some important operations asynchronously.
+type AsyncIndex interface {
+	// Wait will block until asynchronous operations started
+	// before this call have finished or until the specified
+	// timeout has been reached.  If the timeout is reached
+	// an error is returned.
+	Wait(timeout time.Duration) error
+}
+
 type IndexReader interface {
 	TermFieldReader(term []byte, field string) (TermFieldReader, error)
 
diff --git a/test/integration_test.go b/test/integration_test.go
index 374c5f8f..d99f4a3c 100644
--- a/test/integration_test.go
+++ b/test/integration_test.go
@@ -18,8 +18,10 @@ import (
 	"reflect"
 	"regexp"
 	"testing"
+	"time"
 
 	"github.com/blevesearch/bleve"
+	bleveIndex "github.com/blevesearch/bleve/index"
 
 	// we must explicitly include any functionality we plan on testing
 	_ "github.com/blevesearch/bleve/analysis/analyzers/keyword_analyzer"
@@ -28,10 +30,15 @@ import (
 var dataset = flag.String("dataset", "", "only test datasets matching this regex")
 var keepIndex = flag.Bool("keepIndex", false, "keep the index after testing")
 
+var indexType = flag.String("indexType", bleve.Config.DefaultIndexType, "index type to build")
+
 func TestIntegration(t *testing.T) {
 
 	flag.Parse()
 
+	bleve.Config.DefaultIndexType = *indexType
+	t.Logf("using index type %s", *indexType)
+
 	var err error
 	var datasetRegexp *regexp.Regexp
 	if *dataset != "" {
@@ -115,6 +122,19 @@ func runTestDir(t *testing.T, dir string) {
 		}
 	}
 
+	indexInternal, _, err := index.Advanced()
+	if err != nil {
+		t.Fatal(err)
+	}
+	if indexInternal, ok := indexInternal.(bleveIndex.AsyncIndex); ok {
+		start := time.Now()
+		err = indexInternal.Wait(5 * time.Second)
+		if err != nil {
+			t.Fatal(err)
+		}
+		t.Logf("we had to wait for %v", time.Since(start))
+	}
+
 	// read the searches
 	searchBytes, err := ioutil.ReadFile(dir + string(filepath.Separator) + "searches.json")
 	if err != nil {