add some new test utils and test Rows calls on cluster

executor_test.go

@@ -2189,37 +2189,36 @@ func BenchmarkExecutor_Existence_True(b *testing.B)  { benchmarkExistence(true,
 func BenchmarkExecutor_Existence_False(b *testing.B) { benchmarkExistence(false, b) }
 
 func TestExecutor_Execute_Rows(t *testing.T) {
-	c := test.MustRunCluster(t, 1)
+	c := test.MustRunCluster(t, 3)
 	defer c.Close()
-	hldr := test.Holder{Holder: c[0].Server.Holder()}
-	hldr.SetBit("i", "general", 10, 0)
-	hldr.SetBit("i", "general", 10, ShardWidth+1)
-	hldr.SetBit("i", "general", 11, 2)
-	hldr.SetBit("i", "general", 11, ShardWidth+2)
-	hldr.SetBit("i", "general", 12, 2)
-	hldr.SetBit("i", "general", 12, ShardWidth+2)
+	c.CreateField(t, "i", pilosa.IndexOptions{}, "general")
+	c.ImportBits(t, "i", "general", [][2]uint64{
+		{10, 0},
+		{10, ShardWidth + 1},
+		{11, 2},
+		{11, ShardWidth + 2},
+		{12, 2},
+		{12, ShardWidth + 2},
+		{13, 3},
+	})
 
-	if res, err := c[0].API.Query(context.Background(), &pilosa.QueryRequest{Index: "i", Query: `Rows(field=general)`}); err != nil {
-		t.Fatal(err)
-	} else if rows := res.Results[0].(pilosa.RowIdentifiers); !reflect.DeepEqual(rows, pilosa.RowIdentifiers{Rows: []uint64{10, 11, 12}}) {
+	rows := c.Query(t, "i", `Rows(field=general)`).Results[0].(pilosa.RowIdentifiers)
+	if !reflect.DeepEqual(rows, pilosa.RowIdentifiers{Rows: []uint64{10, 11, 12, 13}}) {
 		t.Fatalf("unexpected rows: %+v", rows)
 	}
 
-	if res, err := c[0].API.Query(context.Background(), &pilosa.QueryRequest{Index: "i", Query: `Rows(field=general, limit=2)`}); err != nil {
-		t.Fatal(err)
-	} else if rows := res.Results[0].(pilosa.RowIdentifiers); !reflect.DeepEqual(rows, pilosa.RowIdentifiers{Rows: []uint64{10, 11}}) {
+	rows = c.Query(t, "i", `Rows(field=general, limit=2)`).Results[0].(pilosa.RowIdentifiers)
+	if !reflect.DeepEqual(rows, pilosa.RowIdentifiers{Rows: []uint64{10, 11}}) {
 		t.Fatalf("unexpected rows: %+v", rows)
 	}
 
-	if res, err := c[0].API.Query(context.Background(), &pilosa.QueryRequest{Index: "i", Query: `Rows(field=general, previous=10,limit=2)`}); err != nil {
-		t.Fatal(err)
-	} else if rows := res.Results[0].(pilosa.RowIdentifiers); !reflect.DeepEqual(rows, pilosa.RowIdentifiers{Rows: []uint64{11, 12}}) {
+	rows = c.Query(t, "i", `Rows(field=general, previous=10,limit=2)`).Results[0].(pilosa.RowIdentifiers)
+	if !reflect.DeepEqual(rows, pilosa.RowIdentifiers{Rows: []uint64{11, 12}}) {
 		t.Fatalf("unexpected rows: %+v", rows)
 	}
 
-	if res, err := c[0].API.Query(context.Background(), &pilosa.QueryRequest{Index: "i", Query: `Rows(field=general, column=2)`}); err != nil {
-		t.Fatal(err)
-	} else if rows := res.Results[0].(pilosa.RowIdentifiers); !reflect.DeepEqual(rows, pilosa.RowIdentifiers{Rows: []uint64{11, 12}}) {
+	rows = c.Query(t, "i", `Rows(field=general, column=2)`).Results[0].(pilosa.RowIdentifiers)
+	if !reflect.DeepEqual(rows, pilosa.RowIdentifiers{Rows: []uint64{11, 12}}) {
 		t.Fatalf("unexpected rows: %+v", rows)
 	}
 }

test/pilosa.go

@@ -197,6 +197,74 @@ func (m *Command) RecalculateCaches() error {
 // Cluster represents a Pilosa cluster (multiple Command instances)
 type Cluster []*Command
 
+// Query executes an API.Query through one of the cluster's node's API. It fails
+// the test if there is an error.
+func (c Cluster) Query(t testing.TB, index, query string) pilosa.QueryResponse {
+	if len(c) == 0 {
+		t.Fatal("must have at least one node in cluster to query")
+	}
+
+	return c[0].MustQuery(t, &pilosa.QueryRequest{Index: index, Query: query})
+}
+
+func (c Cluster) ImportBits(t testing.TB, index, field string, rowcols [][2]uint64) {
+	byShard := make(map[uint64][][2]uint64)
+	for _, rowcol := range rowcols {
+		shard := rowcol[1] / pilosa.ShardWidth
+		byShard[shard] = append(byShard[shard], rowcol)
+	}
+
+	for shard, bits := range byShard {
+		rowIDs := make([]uint64, len(bits))
+		colIDs := make([]uint64, len(bits))
+		for i, bit := range bits {
+			rowIDs[i] = bit[0]
+			colIDs[i] = bit[1]
+		}
+		nodes, err := c[0].API.ShardNodes(context.Background(), index, shard)
+		if err != nil {
+			t.Fatalf("getting shard nodes: %v", err)
+		}
+		// TODO won't be necessary to do all nodes once that works hits
+		for _, node := range nodes {
+			for _, com := range c {
+				if com.API.Node().ID != node.ID {
+					continue
+				}
+				err := com.API.Import(context.Background(), &pilosa.ImportRequest{
+					Index:     index,
+					Field:     field,
+					Shard:     shard,
+					RowIDs:    rowIDs,
+					ColumnIDs: colIDs,
+				})
+				if err != nil {
+					t.Fatalf("importing data: %v", err)
+				}
+			}
+		}
+	}
+}
+
+// CreateField creates the index (if necessary) and field specified.
+func (c Cluster) CreateField(t testing.TB, index string, iopts pilosa.IndexOptions, field string, fopts ...pilosa.FieldOption) *pilosa.Field {
+	idx, err := c[0].API.CreateIndex(context.Background(), index, iopts)
+	if err != nil && errors.Cause(err) != pilosa.ErrIndexExists {
+		t.Fatalf("creating index: %v", err)
+	}
+	if idx.Options() != iopts {
+		t.Logf("existing index options:\n%v\ndon't match given opts:\n%v\n in pilosa/test.Cluster.CreateField", idx.Options(), iopts)
+	}
+
+	f, err := c[0].API.CreateField(context.Background(), index, field, fopts...)
+	// we'll assume the field doesn't exist because checking if the options
+	// match seems painful.
+	if err != nil {
+		t.Fatalf("creating field: %v", err)
+	}
+	return f
+}
+
 // Start runs a Cluster
 func (c Cluster) Start() error {
 	var gossipSeeds = make([]string, len(c))