// waitForInitialValue waits for the initial value of
// /keyspaces/test_keyspace/SrvKeyspace to appear, and match the
// provided srvKeyspace.
func waitForInitialValue(t *testing.T, ts topo.Impl, cell string, srvKeyspace *topodatapb.SrvKeyspace) (changes <-chan *topo.WatchData, cancel func()) {
	var current *topo.WatchData
	ctx := context.Background()
	start := time.Now()
	for {
		current, changes, cancel = ts.Watch(ctx, cell, "/keyspaces/test_keyspace/SrvKeyspace")
		if current.Err == topo.ErrNoNode {
			// hasn't appeared yet
			if time.Now().Sub(start) > 10*time.Second {
				t.Fatalf("time out waiting for file to appear")
			}
			time.Sleep(10 * time.Millisecond)
			continue
		}
		if current.Err != nil {
			t.Fatalf("watch failed: %v", current.Err)
		}
		// we got a valid result
		break
	}
	got := &topodatapb.SrvKeyspace{}
	if err := proto.Unmarshal(current.Contents, got); err != nil {
		t.Fatalf("cannot proto-unmarshal data: %v", err)
	}
	if !proto.Equal(got, srvKeyspace) {
		t.Fatalf("got bad data: %v expected: %v", got, srvKeyspace)
	}

	return changes, cancel
}

func checkKeyspaceLockTimeout(ctx context.Context, t *testing.T, ts topo.Impl) {
	lockPath, err := ts.LockKeyspaceForAction(ctx, "test_keyspace", "fake-content")
	if err != nil {
		t.Fatalf("LockKeyspaceForAction: %v", err)
	}

	// test we can't take the lock again
	fastCtx, cancel := context.WithTimeout(ctx, timeUntilLockIsTaken)
	if _, err := ts.LockKeyspaceForAction(fastCtx, "test_keyspace", "unused-fake-content"); err != topo.ErrTimeout {
		t.Fatalf("LockKeyspaceForAction(again): %v", err)
	}
	cancel()

	// test we can interrupt taking the lock
	interruptCtx, cancel := context.WithCancel(ctx)
	go func() {
		time.Sleep(timeUntilLockIsTaken)
		cancel()
	}()
	if _, err := ts.LockKeyspaceForAction(interruptCtx, "test_keyspace", "unused-fake-content"); err != topo.ErrInterrupted {
		t.Fatalf("LockKeyspaceForAction(interrupted): %v", err)
	}

	if err := ts.UnlockKeyspaceForAction(ctx, "test_keyspace", lockPath, "fake-results"); err != nil {
		t.Fatalf("UnlockKeyspaceForAction(): %v", err)
	}

	// test we can't unlock again
	if err := ts.UnlockKeyspaceForAction(ctx, "test_keyspace", lockPath, "fake-results"); err == nil {
		t.Fatalf("UnlockKeyspaceForAction(again) worked")
	}
}

// CheckKeyspaceLock checks we can take a keyspace lock as expected.
func CheckKeyspaceLock(ctx context.Context, t *testing.T, ts topo.Impl) {
	if err := ts.CreateKeyspace(ctx, "test_keyspace", &pb.Keyspace{}); err != nil {
		t.Fatalf("CreateKeyspace: %v", err)
	}

	checkKeyspaceLockTimeout(ctx, t, ts)
	checkKeyspaceLockMissing(ctx, t, ts)
	checkKeyspaceLockUnblocks(ctx, t, ts)
}

func getLocalCell(ctx context.Context, t *testing.T, ts topo.Impl) string {
	cells, err := ts.GetKnownCells(ctx)
	if err != nil {
		t.Fatalf("GetKnownCells: %v", err)
	}
	if len(cells) < 1 {
		t.Fatalf("provided topo.Impl doesn't have enough cells (need at least 1): %v", cells)
	}
	return cells[0]
}

// CopyShardReplications will create the ShardReplication objects in
// the destination topo
func CopyShardReplications(ctx context.Context, fromTS, toTS topo.Impl) {
	keyspaces, err := fromTS.GetKeyspaces(ctx)
	if err != nil {
		log.Fatalf("fromTS.GetKeyspaces: %v", err)
	}

	wg := sync.WaitGroup{}
	rec := concurrency.AllErrorRecorder{}
	for _, keyspace := range keyspaces {
		wg.Add(1)
		go func(keyspace string) {
			defer wg.Done()
			shards, err := fromTS.GetShardNames(ctx, keyspace)
			if err != nil {
				rec.RecordError(fmt.Errorf("GetShardNames(%v): %v", keyspace, err))
				return
			}

			for _, shard := range shards {
				wg.Add(1)
				go func(keyspace, shard string) {
					defer wg.Done()

					// read the source shard to get the cells
					s, _, err := fromTS.GetShard(ctx, keyspace, shard)
					if err != nil {
						rec.RecordError(fmt.Errorf("GetShard(%v, %v): %v", keyspace, shard, err))
						return
					}

					for _, cell := range s.Cells {
						sri, err := fromTS.GetShardReplication(ctx, cell, keyspace, shard)
						if err != nil {
							rec.RecordError(fmt.Errorf("GetShardReplication(%v, %v, %v): %v", cell, keyspace, shard, err))
							continue
						}

						if err := toTS.UpdateShardReplicationFields(ctx, cell, keyspace, shard, func(oldSR *topodatapb.ShardReplication) error {
							*oldSR = *sri.ShardReplication
							return nil
						}); err != nil {
							rec.RecordError(fmt.Errorf("UpdateShardReplicationFields(%v, %v, %v): %v", cell, keyspace, shard, err))
						}
					}
				}(keyspace, shard)
			}
		}(keyspace)
	}
	wg.Wait()
	if rec.HasErrors() {
		log.Fatalf("copyShards failed: %v", rec.Error())
	}
}

// CopyTablets will create the tablets in the destination topo
func CopyTablets(ctx context.Context, fromTS, toTS topo.Impl) {
	cells, err := fromTS.GetKnownCells(ctx)
	if err != nil {
		log.Fatalf("fromTS.GetKnownCells: %v", err)
	}

	wg := sync.WaitGroup{}
	rec := concurrency.AllErrorRecorder{}
	for _, cell := range cells {
		wg.Add(1)
		go func(cell string) {
			defer wg.Done()
			tabletAliases, err := fromTS.GetTabletsByCell(ctx, cell)
			if err != nil {
				rec.RecordError(fmt.Errorf("GetTabletsByCell(%v): %v", cell, err))
			} else {
				for _, tabletAlias := range tabletAliases {
					wg.Add(1)
					go func(tabletAlias *topodatapb.TabletAlias) {
						defer wg.Done()

						// read the source tablet
						tablet, _, err := fromTS.GetTablet(ctx, tabletAlias)
						if err != nil {
							rec.RecordError(fmt.Errorf("GetTablet(%v): %v", tabletAlias, err))
							return
						}

						// try to create the destination
						err = toTS.CreateTablet(ctx, tablet)
						if err == topo.ErrNodeExists {
							// update the destination tablet
							log.Warningf("tablet %v already exists, updating it", tabletAlias)
							_, err = toTS.UpdateTabletFields(ctx, tablet.Alias, func(t *topodatapb.Tablet) error {
								*t = *tablet
								return nil
							})
						}
						if err != nil {
							rec.RecordError(fmt.Errorf("CreateTablet(%v): %v", tabletAlias, err))
							return
						}
					}(tabletAlias)
				}
			}
		}(cell)
	}
	wg.Wait()
	if rec.HasErrors() {
		log.Fatalf("copyTablets failed: %v", rec.Error())
	}
}

// checkWatchInterrupt tests we can interrupt a watch.
func checkWatchInterrupt(t *testing.T, ts topo.Impl) {
	ctx := context.Background()
	cell := getLocalCell(ctx, t, ts)

	// create some data
	srvKeyspace := &topodatapb.SrvKeyspace{
		ShardingColumnName: "user_id",
	}
	if err := ts.UpdateSrvKeyspace(ctx, cell, "test_keyspace", srvKeyspace); err != nil {
		t.Fatalf("UpdateSrvKeyspace(1): %v", err)
	}

	// Start watching, it should work.
	changes, cancel := waitForInitialValue(t, ts, cell, srvKeyspace)

	// Now cancel the watch.
	cancel()

	// Make sure we get the topo.ErrInterrupted notification eventually.
	for {
		wd, ok := <-changes
		if !ok {
			t.Fatalf("watch channel unexpectedly closed")
		}
		if wd.Err == topo.ErrInterrupted {
			// good
			break
		}
		if wd.Err != nil {
			t.Fatalf("bad error returned for deletion: %v", wd.Err)
		}
		// we got something, better be the right value
		got := &topodatapb.SrvKeyspace{}
		if err := proto.Unmarshal(wd.Contents, got); err != nil {
			t.Fatalf("cannot proto-unmarshal data: %v", err)
		}
		if got.ShardingColumnName == "user_id" {
			// good value
			continue
		}
		t.Fatalf("got unknown SrvKeyspace waiting for deletion: %v", got)
	}

	// Now the channel should be closed.
	if wd, ok := <-changes; ok {
		t.Fatalf("got unexpected event after error: %v", wd)
	}

	// And calling cancel() again should just work.
	cancel()
}

// CheckShardLock checks we can take a shard lock
func CheckShardLock(ctx context.Context, t *testing.T, ts topo.Impl) {
	if err := ts.CreateKeyspace(ctx, "test_keyspace", &pb.Keyspace{}); err != nil {
		t.Fatalf("CreateKeyspace: %v", err)
	}
	if err := ts.CreateShard(ctx, "test_keyspace", "10-20", &pb.Shard{
		KeyRange: newKeyRange3("10-20"),
	}); err != nil {
		t.Fatalf("CreateShard: %v", err)
	}

	checkShardLockTimeout(ctx, t, ts)
	checkShardLockMissing(ctx, t, ts)
	checkShardLockUnblocks(ctx, t, ts)
}

func checkListDir(ctx context.Context, t *testing.T, ts topo.Impl, cell string, dirPath string, expected []string) {
	entries, err := ts.ListDir(ctx, cell, dirPath)
	switch err {
	case topo.ErrNoNode:
		if len(expected) != 0 {
			t.Errorf("ListDir(%v) returned ErrNoNode but was expecting %v", dirPath, expected)
		}
	case nil:
		if !reflect.DeepEqual(entries, expected) {
			t.Errorf("ListDir(%v) returned %v but was expecting %v", dirPath, entries, expected)
		}
	default:
		t.Errorf("ListDir(%v) returned unexpected error: %v", dirPath, err)
	}
}

// checkKeyspaceLock checks we can take a keyspace lock as expected.
func checkKeyspaceLock(t *testing.T, ts topo.Impl) {
	ctx := context.Background()
	if err := ts.CreateKeyspace(ctx, "test_keyspace", &topodatapb.Keyspace{}); err != nil {
		t.Fatalf("CreateKeyspace: %v", err)
	}

	t.Log("===      checkKeyspaceLockTimeout")
	checkKeyspaceLockTimeout(ctx, t, ts)

	t.Log("===      checkKeyspaceLockMissing")
	checkKeyspaceLockMissing(ctx, t, ts)

	t.Log("===      checkKeyspaceLockUnblocks")
	checkKeyspaceLockUnblocks(ctx, t, ts)
}

// checkShardLock checks we can take a shard lock
func checkShardLock(t *testing.T, ts topo.Impl) {
	ctx := context.Background()
	if err := ts.CreateKeyspace(ctx, "test_keyspace", &topodatapb.Keyspace{}); err != nil {
		t.Fatalf("CreateKeyspace: %v", err)
	}
	if err := ts.CreateShard(ctx, "test_keyspace", "10-20", &topodatapb.Shard{
		KeyRange: newKeyRange("10-20"),
	}); err != nil {
		t.Fatalf("CreateShard: %v", err)
	}

	t.Log("===     checkShardLockTimeout")
	checkShardLockTimeout(ctx, t, ts)

	t.Log("===     checkShardLockMissing")
	checkShardLockMissing(ctx, t, ts)

	t.Log("===     checkShardLockUnblocks")
	checkShardLockUnblocks(ctx, t, ts)
}

// CopyKeyspaces will create the keyspaces in the destination topo
func CopyKeyspaces(ctx context.Context, fromTS, toTS topo.Impl) {
	keyspaces, err := fromTS.GetKeyspaces(ctx)
	if err != nil {
		log.Fatalf("GetKeyspaces: %v", err)
	}

	wg := sync.WaitGroup{}
	rec := concurrency.AllErrorRecorder{}
	for _, keyspace := range keyspaces {
		wg.Add(1)
		go func(keyspace string) {
			defer wg.Done()

			k, _, err := fromTS.GetKeyspace(ctx, keyspace)
			if err != nil {
				rec.RecordError(fmt.Errorf("GetKeyspace(%v): %v", keyspace, err))
				return
			}

			if err := toTS.CreateKeyspace(ctx, keyspace, k); err != nil {
				if err == topo.ErrNodeExists {
					log.Warningf("keyspace %v already exists", keyspace)
				} else {
					rec.RecordError(fmt.Errorf("CreateKeyspace(%v): %v", keyspace, err))
				}
			}
		}(keyspace)
	}
	wg.Wait()
	if rec.HasErrors() {
		log.Fatalf("copyKeyspaces failed: %v", rec.Error())
	}
}

// checkSrvShardLockUnblocks makes sure that a routine waiting on a lock
// is unblocked when another routine frees the lock
func checkSrvShardLockUnblocks(ctx context.Context, t *testing.T, ts topo.Impl) {
	cell := getLocalCell(ctx, t, ts)
	unblock := make(chan struct{})
	finished := make(chan struct{})

	// as soon as we're unblocked, we try to lock the shard
	go func() {
		<-unblock
		lockPath, err := ts.LockSrvShardForAction(ctx, cell, "test_keyspace", "10-20", "fake-content")
		if err != nil {
			t.Fatalf("LockSrvShardForAction(test, test_keyspace, 10-20) failed: %v", err)
		}
		if err = ts.UnlockSrvShardForAction(ctx, cell, "test_keyspace", "10-20", lockPath, "fake-results"); err != nil {
			t.Fatalf("UnlockSrvShardForAction(test, test_keyspace, 10-20): %v", err)
		}
		close(finished)
	}()

	// lock the shard
	lockPath2, err := ts.LockSrvShardForAction(ctx, cell, "test_keyspace", "10-20", "fake-content")
	if err != nil {
		t.Fatalf("LockSrvShardForAction(test, test_keyspace, 10-20) failed: %v", err)
	}

	// unblock the go routine so it starts waiting
	close(unblock)

	// sleep for a while so we're sure the go routine is blocking
	time.Sleep(timeUntilLockIsTaken)

	if err = ts.UnlockSrvShardForAction(ctx, cell, "test_keyspace", "10-20", lockPath2, "fake-results"); err != nil {
		t.Fatalf("UnlockSrvShardForAction(test, test_keyspace, 10-20): %v", err)
	}

	timeout := time.After(10 * time.Second)
	select {
	case <-finished:
	case <-timeout:
		t.Fatalf("unlocking timed out")
	}
}

// checkSrvVSchema tests the SrvVSchema methods (other than watch).
func checkSrvVSchema(t *testing.T, ts topo.Impl) {
	ctx := context.Background()
	cell := getLocalCell(ctx, t, ts)

	// check GetSrvVSchema returns topo.ErrNoNode if no SrvVSchema
	if _, err := ts.GetSrvVSchema(ctx, cell); err != topo.ErrNoNode {
		t.Errorf("GetSrvVSchema(not set): %v", err)
	}

	srvVSchema := &vschemapb.SrvVSchema{
		Keyspaces: map[string]*vschemapb.Keyspace{
			"test_keyspace": {
				Sharded: true,
			},
		},
	}
	if err := ts.UpdateSrvVSchema(ctx, cell, srvVSchema); err != nil {
		t.Errorf("UpdateSrvVSchema(1): %v", err)
	}
	if v, err := ts.GetSrvVSchema(ctx, cell); err != nil || !proto.Equal(srvVSchema, v) {
		t.Errorf("GetSrvVSchema(valid): %v %v", err, v)
	}
}

func checkShardLockMissing(ctx context.Context, t *testing.T, ts topo.Impl) {
	// test we can't lock a non-existing shard
	if _, err := ts.LockShardForAction(ctx, "test_keyspace", "20-30", "fake-content"); err == nil {
		t.Fatalf("LockShardForAction(test_keyspace/20-30) worked for non-existing shard")
	}
}

// TopoServerTestSuite runs the full topo.Impl test suite.
// The factory method should return a topo server that has a single cell
// called 'test'.
func TopoServerTestSuite(t *testing.T, factory func() topo.Impl) {
	var ts topo.Impl

	t.Log("=== checkKeyspace")
	ts = factory()
	checkKeyspace(t, ts)
	ts.Close()

	t.Log("=== checkShard")
	ts = factory()
	checkShard(t, ts)
	ts.Close()

	t.Log("=== checkTablet")
	ts = factory()
	checkTablet(t, ts)
	ts.Close()

	t.Log("=== checkShardReplication")
	ts = factory()
	checkShardReplication(t, ts)
	ts.Close()

	t.Log("=== checkSrvKeyspace")
	ts = factory()
	checkSrvKeyspace(t, ts)
	ts.Close()

	t.Log("=== checkWatchSrvKeyspace")
	ts = factory()
	checkWatchSrvKeyspace(t, ts)
	ts.Close()

	t.Log("=== checkSrvVSchema")
	ts = factory()
	checkSrvVSchema(t, ts)
	ts.Close()

	t.Log("=== checkWatchSrvVSchema")
	ts = factory()
	checkWatchSrvVSchema(t, ts)
	ts.Close()

	t.Log("=== checkKeyspaceLock")
	ts = factory()
	checkKeyspaceLock(t, ts)
	ts.Close()

	t.Log("=== checkShardLock")
	ts = factory()
	checkShardLock(t, ts)
	ts.Close()

	t.Log("=== checkVSchema")
	ts = factory()
	checkVSchema(t, ts)
	ts.Close()
}

// CopyShards will create the shards in the destination topo
func CopyShards(ctx context.Context, fromTS, toTS topo.Impl, deleteKeyspaceShards bool) {
	keyspaces, err := fromTS.GetKeyspaces(ctx)
	if err != nil {
		log.Fatalf("fromTS.GetKeyspaces: %v", err)
	}

	wg := sync.WaitGroup{}
	rec := concurrency.AllErrorRecorder{}
	for _, keyspace := range keyspaces {
		wg.Add(1)
		go func(keyspace string) {
			defer wg.Done()
			shards, err := fromTS.GetShardNames(ctx, keyspace)
			if err != nil {
				rec.RecordError(fmt.Errorf("GetShardNames(%v): %v", keyspace, err))
				return
			}

			if deleteKeyspaceShards {
				if err := toTS.DeleteKeyspaceShards(ctx, keyspace); err != nil {
					rec.RecordError(fmt.Errorf("DeleteKeyspaceShards(%v): %v", keyspace, err))
					return
				}
			}

			for _, shard := range shards {
				wg.Add(1)
				go func(keyspace, shard string) {
					defer wg.Done()
					if err := toTS.CreateShard(ctx, keyspace, shard, &topodatapb.Shard{}); err != nil {
						if err == topo.ErrNodeExists {
							log.Warningf("shard %v/%v already exists", keyspace, shard)
						} else {
							rec.RecordError(fmt.Errorf("CreateShard(%v, %v): %v", keyspace, shard, err))
							return
						}
					}

					s, _, err := fromTS.GetShard(ctx, keyspace, shard)
					if err != nil {
						rec.RecordError(fmt.Errorf("GetShard(%v, %v): %v", keyspace, shard, err))
						return
					}

					_, toV, err := toTS.GetShard(ctx, keyspace, shard)
					if err != nil {
						rec.RecordError(fmt.Errorf("toTS.GetShard(%v, %v): %v", keyspace, shard, err))
						return
					}

					if _, err := toTS.UpdateShard(ctx, keyspace, shard, s, toV); err != nil {
						rec.RecordError(fmt.Errorf("UpdateShard(%v, %v): %v", keyspace, shard, err))
					}
				}(keyspace, shard)
			}
		}(keyspace)
	}
	wg.Wait()
	if rec.HasErrors() {
		log.Fatalf("copyShards failed: %v", rec.Error())
	}
}

// checkSrvKeyspace tests the SrvKeyspace methods (other than watch).
func checkSrvKeyspace(t *testing.T, ts topo.Impl) {
	ctx := context.Background()
	cell := getLocalCell(ctx, t, ts)

	// test cell/keyspace entries (SrvKeyspace)
	srvKeyspace := &topodatapb.SrvKeyspace{
		Partitions: []*topodatapb.SrvKeyspace_KeyspacePartition{
			{
				ServedType: topodatapb.TabletType_MASTER,
				ShardReferences: []*topodatapb.ShardReference{
					{
						Name: "-80",
						KeyRange: &topodatapb.KeyRange{
							End: []byte{0x80},
						},
					},
				},
			},
		},
		ShardingColumnName: "video_id",
		ShardingColumnType: topodatapb.KeyspaceIdType_UINT64,
		ServedFrom: []*topodatapb.SrvKeyspace_ServedFrom{
			{
				TabletType: topodatapb.TabletType_REPLICA,
				Keyspace:   "other_keyspace",
			},
		},
	}
	if err := ts.UpdateSrvKeyspace(ctx, cell, "test_keyspace", srvKeyspace); err != nil {
		t.Errorf("UpdateSrvKeyspace(1): %v", err)
	}
	if _, err := ts.GetSrvKeyspace(ctx, cell, "test_keyspace666"); err != topo.ErrNoNode {
		t.Errorf("GetSrvKeyspace(invalid): %v", err)
	}
	if k, err := ts.GetSrvKeyspace(ctx, cell, "test_keyspace"); err != nil || !proto.Equal(srvKeyspace, k) {
		t.Errorf("GetSrvKeyspace(valid): %v %v", err, k)
	}
	if k, err := ts.GetSrvKeyspaceNames(ctx, cell); err != nil || len(k) != 1 || k[0] != "test_keyspace" {
		t.Errorf("GetSrvKeyspaceNames(): %v", err)
	}

	// check that updating a SrvKeyspace out of the blue works
	if err := ts.UpdateSrvKeyspace(ctx, cell, "unknown_keyspace_so_far", srvKeyspace); err != nil {
		t.Fatalf("UpdateSrvKeyspace(2): %v", err)
	}
	if k, err := ts.GetSrvKeyspace(ctx, cell, "unknown_keyspace_so_far"); err != nil || !proto.Equal(srvKeyspace, k) {
		t.Errorf("GetSrvKeyspace(out of the blue): %v %v", err, *k)
	}

	// Delete the SrvKeyspace.
	if err := ts.DeleteSrvKeyspace(ctx, cell, "unknown_keyspace_so_far"); err != nil {
		t.Fatalf("DeleteSrvKeyspace: %v", err)
	}
	if _, err := ts.GetSrvKeyspace(ctx, cell, "unknown_keyspace_so_far"); err != topo.ErrNoNode {
		t.Errorf("GetSrvKeyspace(deleted) got %v, want ErrNoNode", err)
	}
}

func checkSrvShardLockGeneral(ctx context.Context, t *testing.T, ts topo.Impl) {
	cell := getLocalCell(ctx, t, ts)

	// make sure we can create the lock even if no directory exists
	lockPath, err := ts.LockSrvShardForAction(ctx, cell, "test_keyspace", "10-20", "fake-content")
	if err != nil {
		t.Fatalf("LockSrvShardForAction: %v", err)
	}

	if err := ts.UnlockSrvShardForAction(ctx, cell, "test_keyspace", "10-20", lockPath, "fake-results"); err != nil {
		t.Fatalf("UnlockShardForAction: %v", err)
	}

	// now take the lock again after the root exists
	lockPath, err = ts.LockSrvShardForAction(ctx, cell, "test_keyspace", "10-20", "fake-content")
	if err != nil {
		t.Fatalf("LockSrvShardForAction: %v", err)
	}

	// test we can't take the lock again
	fastCtx, cancel := context.WithTimeout(ctx, timeUntilLockIsTaken)
	if _, err := ts.LockSrvShardForAction(fastCtx, cell, "test_keyspace", "10-20", "unused-fake-content"); err != topo.ErrTimeout {
		t.Fatalf("LockSrvShardForAction(again): %v", err)
	}
	cancel()

	// test we can interrupt taking the lock
	interruptCtx, cancel := context.WithCancel(ctx)
	go func() {
		time.Sleep(timeUntilLockIsTaken)
		cancel()
	}()
	if _, err := ts.LockSrvShardForAction(interruptCtx, cell, "test_keyspace", "10-20", "unused-fake-content"); err != topo.ErrInterrupted {
		t.Fatalf("LockSrvShardForAction(interrupted): %v", err)
	}

	// unlock now
	if err := ts.UnlockSrvShardForAction(ctx, cell, "test_keyspace", "10-20", lockPath, "fake-results"); err != nil {
		t.Fatalf("UnlockSrvShardForAction(): %v", err)
	}

	// test we can't unlock again
	if err := ts.UnlockSrvShardForAction(ctx, cell, "test_keyspace", "10-20", lockPath, "fake-results"); err == nil {
		t.Error("UnlockSrvShardForAction(again) worked")
	}

}

// checkShard verifies the Shard operations work correctly
func checkShard(t *testing.T, ts topo.Impl) {
	ctx := context.Background()
	tts := topo.Server{Impl: ts}

	if err := ts.CreateKeyspace(ctx, "test_keyspace", &topodatapb.Keyspace{}); err != nil {
		t.Fatalf("CreateKeyspace: %v", err)
	}

	shard := &topodatapb.Shard{
		KeyRange: newKeyRange("b0-c0"),
	}
	if err := ts.CreateShard(ctx, "test_keyspace", "b0-c0", shard); err != nil {
		t.Fatalf("CreateShard: %v", err)
	}
	if err := ts.CreateShard(ctx, "test_keyspace", "b0-c0", shard); err != topo.ErrNodeExists {
		t.Errorf("CreateShard called second time, got: %v", err)
	}

	// Delete shard and see if we can re-create it.
	if err := ts.DeleteShard(ctx, "test_keyspace", "b0-c0"); err != nil {
		t.Fatalf("DeleteShard: %v", err)
	}
	if err := ts.CreateShard(ctx, "test_keyspace", "b0-c0", shard); err != nil {
		t.Fatalf("CreateShard: %v", err)
	}

	// Delete ALL shards.
	if err := ts.DeleteKeyspaceShards(ctx, "test_keyspace"); err != nil {
		t.Fatalf("DeleteKeyspaceShards: %v", err)
	}
	if err := ts.CreateShard(ctx, "test_keyspace", "b0-c0", shard); err != nil {
		t.Fatalf("CreateShard: %v", err)
	}

	if _, _, err := ts.GetShard(ctx, "test_keyspace", "666"); err != topo.ErrNoNode {
		t.Errorf("GetShard(666): %v", err)
	}

	shard, version, err := ts.GetShard(ctx, "test_keyspace", "b0-c0")
	if err != nil {
		t.Errorf("GetShard: %v", err)
	}
	if want := newKeyRange("b0-c0"); !key.KeyRangeEqual(shard.KeyRange, want) {
		t.Errorf("shard.KeyRange: want %v, got %v", want, shard.KeyRange)
	}
	master := &topodatapb.TabletAlias{Cell: "ny", Uid: 1}
	shard.MasterAlias = master
	shard.KeyRange = newKeyRange("b0-c0")
	shard.ServedTypes = []*topodatapb.Shard_ServedType{
		{
			TabletType: topodatapb.TabletType_MASTER,
		},
		{
			TabletType: topodatapb.TabletType_REPLICA,
			Cells:      []string{"c1"},
		},
		{
			TabletType: topodatapb.TabletType_RDONLY,
		},
	}
	shard.SourceShards = []*topodatapb.Shard_SourceShard{
		{
			Uid:      1,
			Keyspace: "source_ks",
			Shard:    "b8-c0",
			KeyRange: newKeyRange("b8-c0"),
			Tables:   []string{"table1", "table2"},
		},
	}
	shard.TabletControls = []*topodatapb.Shard_TabletControl{
		{
			TabletType:        topodatapb.TabletType_MASTER,
			Cells:             []string{"c1", "c2"},
			BlacklistedTables: []string{"black1", "black2"},
		},
		{
			TabletType:          topodatapb.TabletType_REPLICA,
			DisableQueryService: true,
		},
	}
	if _, err := ts.UpdateShard(ctx, "test_keyspace", "b0-c0", shard, version); err != nil {
		t.Errorf("UpdateShard: %v", err)
	}

	other := &topodatapb.TabletAlias{Cell: "ny", Uid: 82873}
	_, err = tts.UpdateShardFields(ctx, "test_keyspace", "b0-c0", func(si *topo.ShardInfo) error {
		si.MasterAlias = other
		return nil
	})
	if err != nil {
		t.Fatalf("UpdateShardFields error: %v", err)
	}

	s, _, err := ts.GetShard(ctx, "test_keyspace", "b0-c0")
	if err != nil {
		t.Fatalf("GetShard: %v", err)
	}
	if *s.MasterAlias != *other {
		t.Fatalf("shard.MasterAlias = %v, want %v", s.MasterAlias, other)
//.........这里部分代码省略.........