use existing func, add more unit tests to cover all errors

Signed-off-by: deepthi <deepthi@planetscale.com>
vitessio · Jul 17, 2020 · 2d9c8e3 · 2d9c8e3
1 parent 166607d
commit 2d9c8e3
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 27 deletions.
diff --git a/go/vt/discovery/tablet_picker.go b/go/vt/discovery/tablet_picker.go
@@ -58,11 +58,13 @@ func NewTabletPicker(ts *topo.Server, cells []string, keyspace, shard, tabletTyp
 	}, nil
 }
 
-// PickForStreaming picks all healthy tablets including the non-serving ones.
+// PickForStreaming picks an available tablet
+// All tablets that belong to tp.cells are evaluated and one is
+// chosen at random
 func (tp *TabletPicker) PickForStreaming(ctx context.Context) (*topodatapb.Tablet, error) {
 	candidates := tp.getAllTablets(ctx)
 	if len(candidates) == 0 {
-		return nil, vterrors.Errorf(vtrpcpb.Code_NOT_FOUND, "no tablets available for %v %v %v", tp.cells, tp.keyspace, tp.shard)
+		return nil, vterrors.Errorf(vtrpcpb.Code_NOT_FOUND, "no tablets available for cells:%v, keyspace/shard:%v/%v, tablet types:%v", tp.cells, tp.keyspace, tp.shard, tp.tabletTypes)
 	}
 	for {
 		idx := 0
@@ -81,7 +83,7 @@ func (tp *TabletPicker) PickForStreaming(ctx context.Context) (*topodatapb.Table
 			}
 			continue
 		}
-		if !tabletTypeIn(ti.Tablet.Type, tp.tabletTypes) {
+		if !topoproto.IsTypeInList(ti.Tablet.Type, tp.tabletTypes) {
 			// tablet is not of one of the desired types
 			continue
 		}
@@ -99,17 +101,9 @@ func (tp *TabletPicker) PickForStreaming(ctx context.Context) (*topodatapb.Table
 		_ = conn.Close(ctx)
 		return ti.Tablet, nil
 	}
-	return nil, vterrors.Errorf(vtrpcpb.Code_NOT_FOUND, "can't find any healthy source tablet for %v %v %v", tp.keyspace, tp.shard, tp.tabletTypes)
+	return nil, vterrors.Errorf(vtrpcpb.Code_NOT_FOUND, "can't find any healthy source tablet for keyspace/shard:%v/%v tablet types:%v", tp.keyspace, tp.shard, tp.tabletTypes)
 }
 
-func tabletTypeIn(tabletType topodatapb.TabletType, tabletTypes []topodatapb.TabletType) bool {
-	for _, t := range tabletTypes {
-		if tabletType == t {
-			return true
-		}
-	}
-	return false
-}
 func (tp *TabletPicker) getAllTablets(ctx context.Context) []*topodatapb.TabletAlias {
 	// Special handling for MASTER tablet type
 	// Since there is only one master, we ignore cell and find the master

diff --git a/go/vt/discovery/tablet_picker_test.go b/go/vt/discovery/tablet_picker_test.go
@@ -277,13 +277,18 @@ func TestPickUsingCellAlias(t *testing.T) {
 
 func TestPickError(t *testing.T) {
 	te := newPickerTestEnv(t, []string{"cell"})
-	defer deleteTablet(te, addTablet(te, 100, topodatapb.TabletType_REPLICA, "cell", false, false))
-
 	_, err := NewTabletPicker(te.topoServ, te.cells, te.keyspace, te.shard, "badtype")
 	assert.EqualError(t, err, "failed to parse list of tablet types: badtype")
 
-	_, err = NewTabletPicker(te.topoServ, te.cells, te.keyspace, te.shard, "replica,rdonly")
+	tp, err := NewTabletPicker(te.topoServ, te.cells, te.keyspace, te.shard, "replica,rdonly")
 	require.NoError(t, err)
+	ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond)
+	defer cancel()
+	_, err = tp.PickForStreaming(ctx)
+	require.EqualError(t, err, "no tablets available for cells:[cell], keyspace/shard:ks/0, tablet types:[REPLICA RDONLY]")
+	defer deleteTablet(te, addTablet(te, 200, topodatapb.TabletType_REPLICA, "cell", false, false))
+	_, err = tp.PickForStreaming(ctx)
+	require.EqualError(t, err, "can't find any healthy source tablet for keyspace/shard:ks/0 tablet types:[REPLICA RDONLY]")
 }
 
 type pickerTestEnv struct {
@@ -334,19 +339,17 @@ func addTablet(te *pickerTestEnv, id int, tabletType topodatapb.TabletType, cell
 	err := te.topoServ.CreateTablet(context.Background(), tablet)
 	require.NoError(te.t, err)
 
-	var herr string
-	if !healthy {
-		herr = "err"
+	if healthy {
+		_ = createFixedHealthConn(tablet, &querypb.StreamHealthResponse{
+			Serving: serving,
+			Target: &querypb.Target{
+				Keyspace:   te.keyspace,
+				Shard:      te.shard,
+				TabletType: tabletType,
+			},
+			RealtimeStats: &querypb.RealtimeStats{HealthError: ""},
+		})
 	}
-	_ = createFixedHealthConn(tablet, &querypb.StreamHealthResponse{
-		Serving: serving,
-		Target: &querypb.Target{
-			Keyspace:   te.keyspace,
-			Shard:      te.shard,
-			TabletType: tabletType,
-		},
-		RealtimeStats: &querypb.RealtimeStats{HealthError: herr},
-	})
 
 	return tablet
 }