package managerserver

import (
	"context"
	"encoding/json"
	"errors"
	"net/http"
	"net/http/httptest"
	"sync"
	"testing"

	"github.com/pingcap/kvproto/pkg/keyspacepb"
	"github.com/tidbcloud/aws-shared-provider/cmd/manager/app/config"
	"github.com/tidbcloud/aws-shared-provider/pkg/controllers/keyspacepoolconfig"
	"github.com/tidbcloud/aws-shared-provider/pkg/keyspacemigration"
	wrapmanager "github.com/tidbcloud/aws-shared-provider/pkg/manager"
	"github.com/tidbcloud/aws-shared-provider/pkg/scheme"
	"github.com/tidbcloud/aws-shared-provider/pkg/tidbmanager"
	"github.com/tidbcloud/aws-shared-provider/pkg/types"
	utilslabels "github.com/tidbcloud/aws-shared-provider/pkg/utils/labels"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils/pd"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/client/fake"
	ctrlmanager "sigs.k8s.io/controller-runtime/pkg/manager"
)

// fakeWrapManager is a minimal wrapmanager.Manager that only exposes a client.
// The embedded ctrlmanager.Manager is left as a nil interface, so any method
// we don't override (e.g. Add, Start, GetCache) will panic. Tests must only
// exercise code paths that touch GetClient.
type fakeWrapManager struct {
	ctrlmanager.Manager
	cli client.Client
}

func (f *fakeWrapManager) GetClient() client.Client { return f.cli }

var _ wrapmanager.Manager = (*fakeWrapManager)(nil)

// fakePDClient is a hand-rolled stub for pd.HttpPdClient.
type fakePDClient struct {
	keyspaces            map[string]*pd.KeyspaceMeta
	getKeyspaceErr       error
	updateKeyspaceErr    error
	updateKeyspaceCalled int
	lastUpdateConfig     *pd.KeyspaceConfig
	metaServiceGroups    *[]pd.MetaServiceGroup
	getGroupListErr      error
	mu                   sync.Mutex
}

func (f *fakePDClient) GetClusterID(_ context.Context) (uint64, error) { return 0, nil }
func (f *fakePDClient) GetKeyspaceByID(_ context.Context, _ uint32) (*pd.KeyspaceMeta, error) {
	return nil, nil
}
func (f *fakePDClient) GetKeyspaceByName(_ context.Context, name string) (*pd.KeyspaceMeta, error) {
	if f.getKeyspaceErr != nil {
		return nil, f.getKeyspaceErr
	}
	ks, ok := f.keyspaces[name]
	if !ok {
		return nil, errors.New("keyspace does not exist")
	}
	return ks, nil
}
func (f *fakePDClient) CreateKeyspace(_ context.Context, _ *pd.KeyspaceMeta) (*pd.KeyspaceMeta, error) {
	return nil, nil
}
func (f *fakePDClient) UpdateKeyspaceState(_ context.Context, _ string, _ keyspacepb.KeyspaceState) error {
	return nil
}
func (f *fakePDClient) UpdateKeyspaceConfig(_ context.Context, name string, c *pd.KeyspaceConfig) error {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.updateKeyspaceCalled++
	f.lastUpdateConfig = c
	if f.updateKeyspaceErr != nil {
		return f.updateKeyspaceErr
	}
	if ks := f.keyspaces[name]; ks != nil {
		ks.Config.MetaServiceGroupID = c.MetaServiceGroupID
	}
	return nil
}
func (f *fakePDClient) PostResourceGroup(_ context.Context, _ *pd.ResourceGroup) error { return nil }
func (f *fakePDClient) PutResourceGroup(_ context.Context, _ *pd.ResourceGroup) error  { return nil }
func (f *fakePDClient) DeleteResourceGroup(_ context.Context, _ string) error          { return nil }
func (f *fakePDClient) SetKeyspaceRuVersion(_ context.Context, _ string, _ *pd.SetKeyspaceRuVersionRequest) error {
	return nil
}
func (f *fakePDClient) GetMetaServiceGroupList(_ context.Context) (*[]pd.MetaServiceGroup, error) {
	return f.metaServiceGroups, f.getGroupListErr
}

// fakeTidbWorkerApi is a hand-rolled stub for tidbworker.Api.
type fakeTidbWorkerApi struct {
	mu           sync.Mutex
	pauseErr     error
	resumeErr    error
	pauseCalled  int
	resumeCalled int
	pausedAt     int // sequence number when pause was called
	resumedAt    int
	callSeq      int
}

func (f *fakeTidbWorkerApi) Clean(_ context.Context, _ string) error { return nil }
func (f *fakeTidbWorkerApi) Pause(_ context.Context, _ string) error {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.callSeq++
	f.pauseCalled++
	f.pausedAt = f.callSeq
	return f.pauseErr
}
func (f *fakeTidbWorkerApi) Resume(_ context.Context, _ string) error {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.callSeq++
	f.resumeCalled++
	f.resumedAt = f.callSeq
	return f.resumeErr
}

// testTiDBWorkerNS is the namespace the test handler looks in for tidb-worker
// pods. We use a distinct value (vs. the production default) to make sure the
// handler honours s.cfg.TiDBWorkerNS rather than a hard-coded constant.
const testTiDBWorkerNS = "tidb-worker-test"

// newServerForTest builds a server wired with the supplied dependencies.
// Callers can leave any field nil to get default empty behavior.
func newServerForTest(
	pdcli pd.HttpPdClient,
	worker *fakeTidbWorkerApi,
	locker keyspacemigration.Locker,
	tManagers map[string]tidbmanager.TidbManager,
	objs ...client.Object,
) *server {
	cli := fake.NewClientBuilder().WithScheme(scheme.Scheme).WithObjects(objs...).Build()
	if locker == nil {
		locker = keyspacemigration.New()
	}
	if tManagers == nil {
		tManagers = map[string]tidbmanager.TidbManager{}
	}
	return &server{
		manager:         &fakeWrapManager{cli: cli},
		tManagers:       tManagers,
		pdClient:        pdcli,
		tidbWorkerApi:   worker,
		migrationLocker: locker,
		createRateStore: keyspacepoolconfig.NewCreateRateStore(0),
		cfg:             &config.Config{TiDBWorkerNS: testTiDBWorkerNS},
	}
}

func TestMigrateMetaServiceGroup_MissingKeyspaceParam(t *testing.T) {
	s := newServerForTest(&fakePDClient{}, &fakeTidbWorkerApi{}, nil, nil)

	rec := httptest.NewRecorder()
	req := httptest.NewRequest(http.MethodPost,
		"/api/keyspace/migrate_meta_service_group?meta-service-group-id=g1", nil)
	s.migrateMetaServiceGroup(rec, req)

	if rec.Code != http.StatusBadRequest {
		t.Fatalf("expected 400, got %d body=%s", rec.Code, rec.Body.String())
	}
}

func TestMigrateMetaServiceGroup_MissingGroupParam(t *testing.T) {
	s := newServerForTest(&fakePDClient{}, &fakeTidbWorkerApi{}, nil, nil)

	rec := httptest.NewRecorder()
	req := httptest.NewRequest(http.MethodPost,
		"/api/keyspace/migrate_meta_service_group?keyspace-name=ks", nil)
	s.migrateMetaServiceGroup(rec, req)

	if rec.Code != http.StatusBadRequest {
		t.Fatalf("expected 400, got %d body=%s", rec.Code, rec.Body.String())
	}
}

func TestMigrateMetaServiceGroup_KeyspaceNotFound(t *testing.T) {
	pdcli := &fakePDClient{
		keyspaces: map[string]*pd.KeyspaceMeta{},
	}
	s := newServerForTest(pdcli, &fakeTidbWorkerApi{}, nil, nil)

	rec := httptest.NewRecorder()
	req := httptest.NewRequest(http.MethodPost,
		"/api/keyspace/migrate_meta_service_group?keyspace-name=ks-missing&meta-service-group-id=g1", nil)
	s.migrateMetaServiceGroup(rec, req)

	if rec.Code != http.StatusBadRequest {
		t.Fatalf("expected 400, got %d body=%s", rec.Code, rec.Body.String())
	}
}

func TestMigrateMetaServiceGroup_NoopWhenAlreadyOnTargetGroup(t *testing.T) {
	pdcli := &fakePDClient{
		keyspaces: map[string]*pd.KeyspaceMeta{
			"ks-a": {Name: "ks-a", Config: pd.KeyspaceConfig{MetaServiceGroupID: "g1"}},
		},
	}
	worker := &fakeTidbWorkerApi{}
	s := newServerForTest(pdcli, worker, nil, nil)

	rec := httptest.NewRecorder()
	req := httptest.NewRequest(http.MethodPost,
		"/api/keyspace/migrate_meta_service_group?keyspace-name=ks-a&meta-service-group-id=g1", nil)
	s.migrateMetaServiceGroup(rec, req)

	if rec.Code != http.StatusOK {
		t.Fatalf("expected 200, got %d body=%s", rec.Code, rec.Body.String())
	}
	if pdcli.updateKeyspaceCalled != 0 {
		t.Fatalf("noop must not update pd; called %d times", pdcli.updateKeyspaceCalled)
	}
	if worker.pauseCalled != 0 {
		t.Fatalf("noop must not pause tidb-worker; called %d times", worker.pauseCalled)
	}

	var got migrateResult
	if err := json.Unmarshal(rec.Body.Bytes(), &got); err != nil {
		t.Fatalf("decode response: %v", err)
	}
	if got.Status != "noop" {
		t.Fatalf("want status=noop, got %q", got.Status)
	}
}

func TestMigrateMetaServiceGroup_TargetGroupMustExist(t *testing.T) {
	pdcli := &fakePDClient{
		keyspaces: map[string]*pd.KeyspaceMeta{
			"ks-a": {Name: "ks-a"},
		},
		metaServiceGroups: &[]pd.MetaServiceGroup{{ID: "g1"}},
	}
	s := newServerForTest(pdcli, &fakeTidbWorkerApi{}, nil, nil)

	rec := httptest.NewRecorder()
	req := httptest.NewRequest(http.MethodPost,
		"/api/keyspace/migrate_meta_service_group?keyspace-name=ks-a&meta-service-group-id=g-unknown", nil)
	s.migrateMetaServiceGroup(rec, req)

	if rec.Code != http.StatusBadRequest {
		t.Fatalf("expected 400, got %d body=%s", rec.Code, rec.Body.String())
	}
}

func TestMigrateMetaServiceGroup_BlockedByRunningTidb(t *testing.T) {
	pdcli := &fakePDClient{
		keyspaces: map[string]*pd.KeyspaceMeta{
			"ks-a": {Name: "ks-a"},
		},
		metaServiceGroups: &[]pd.MetaServiceGroup{{ID: "g1"}},
	}
	tManagers := map[string]tidbmanager.TidbManager{
		"tidb-pool": &testTidbManager{
			active: map[string]testWalkInfo{
				"ks-a": {endpoints: []string{"pod-x,10.0.0.1"}, replicas: 1},
			},
		},
	}
	worker := &fakeTidbWorkerApi{}
	s := newServerForTest(pdcli, worker, nil, tManagers)

	rec := httptest.NewRecorder()
	req := httptest.NewRequest(http.MethodPost,
		"/api/keyspace/migrate_meta_service_group?keyspace-name=ks-a&meta-service-group-id=g1", nil)
	s.migrateMetaServiceGroup(rec, req)

	if rec.Code != http.StatusConflict {
		t.Fatalf("expected 409, got %d body=%s", rec.Code, rec.Body.String())
	}
	if pdcli.updateKeyspaceCalled != 0 {
		t.Fatalf("must not update pd when tidb is running; called %d times", pdcli.updateKeyspaceCalled)
	}
	if worker.pauseCalled != 0 {
		t.Fatalf("must not pause tidb-worker when pre-check fails; called %d times", worker.pauseCalled)
	}
}

func TestMigrateMetaServiceGroup_BlockedByLightningPod(t *testing.T) {
	pdcli := &fakePDClient{
		keyspaces: map[string]*pd.KeyspaceMeta{
			"ks-a": {Name: "ks-a"},
		},
		metaServiceGroups: &[]pd.MetaServiceGroup{{ID: "g1"}},
	}
	lightningPod := &corev1.Pod{
		ObjectMeta: metav1.ObjectMeta{
			Name:      "lightning-x",
			Namespace: importNamespace,
			Labels:    map[string]string{types.KeyspaceNameLabel: "ks-a"},
		},
	}
	worker := &fakeTidbWorkerApi{}
	s := newServerForTest(pdcli, worker, nil, nil, lightningPod)

	rec := httptest.NewRecorder()
	req := httptest.NewRequest(http.MethodPost,
		"/api/keyspace/migrate_meta_service_group?keyspace-name=ks-a&meta-service-group-id=g1", nil)
	s.migrateMetaServiceGroup(rec, req)

	if rec.Code != http.StatusConflict {
		t.Fatalf("expected 409, got %d body=%s", rec.Code, rec.Body.String())
	}
}

func TestMigrateMetaServiceGroup_BlockedByDumplingPod(t *testing.T) {
	pdcli := &fakePDClient{
		keyspaces: map[string]*pd.KeyspaceMeta{
			"ks-a": {Name: "ks-a"},
		},
		metaServiceGroups: &[]pd.MetaServiceGroup{{ID: "g1"}},
	}
	dumplingPod := &corev1.Pod{
		ObjectMeta: metav1.ObjectMeta{
			Name:      "dumpling-x",
			Namespace: exportNamespace,
			Labels:    map[string]string{types.KeyspaceNameLabel: "ks-a"},
		},
	}
	worker := &fakeTidbWorkerApi{}
	s := newServerForTest(pdcli, worker, nil, nil, dumplingPod)

	rec := httptest.NewRecorder()
	req := httptest.NewRequest(http.MethodPost,
		"/api/keyspace/migrate_meta_service_group?keyspace-name=ks-a&meta-service-group-id=g1", nil)
	s.migrateMetaServiceGroup(rec, req)

	if rec.Code != http.StatusConflict {
		t.Fatalf("expected 409, got %d body=%s", rec.Code, rec.Body.String())
	}
}

func TestMigrateMetaServiceGroup_BlockedByTidbWorkerPod(t *testing.T) {
	pdcli := &fakePDClient{
		keyspaces: map[string]*pd.KeyspaceMeta{
			"ks-a": {Name: "ks-a"},
		},
		metaServiceGroups: &[]pd.MetaServiceGroup{{ID: "g1"}},
	}
	// A live tidb-worker pod for ks-a in the configured tidb-worker
	// namespace, with the same labels the scaler stamps.
	workerPod := &corev1.Pod{
		ObjectMeta: metav1.ObjectMeta{
			Name:      "tidb-worker-shared-1-0",
			Namespace: testTiDBWorkerNS,
			Labels: map[string]string{
				types.KeyspaceNameLabelKey:    "ks-a",
				utilslabels.ComponentLabelKey: "worker",
			},
		},
	}
	worker := &fakeTidbWorkerApi{}
	s := newServerForTest(pdcli, worker, nil, nil, workerPod)

	rec := httptest.NewRecorder()
	req := httptest.NewRequest(http.MethodPost,
		"/api/keyspace/migrate_meta_service_group?keyspace-name=ks-a&meta-service-group-id=g1", nil)
	s.migrateMetaServiceGroup(rec, req)

	if rec.Code != http.StatusConflict {
		t.Fatalf("expected 409, got %d body=%s", rec.Code, rec.Body.String())
	}
	if worker.pauseCalled != 0 {
		t.Fatalf("must not pause when pre-check fails; called %d times", worker.pauseCalled)
	}
}

// A pod in the tidb-worker namespace that does not carry the keyspace label
// must be ignored by the check.
func TestMigrateMetaServiceGroup_OtherKeyspacesTidbWorkerPodIgnored(t *testing.T) {
	pdcli := &fakePDClient{
		keyspaces: map[string]*pd.KeyspaceMeta{
			"ks-a": {Name: "ks-a", Config: pd.KeyspaceConfig{MetaServiceGroupID: "g0"}},
		},
		metaServiceGroups: &[]pd.MetaServiceGroup{{ID: "g0"}, {ID: "g1"}},
	}
	otherKsWorker := &corev1.Pod{
		ObjectMeta: metav1.ObjectMeta{
			Name:      "tidb-worker-shared-other-0",
			Namespace: testTiDBWorkerNS,
			Labels: map[string]string{
				types.KeyspaceNameLabelKey:    "ks-other",
				utilslabels.ComponentLabelKey: "worker",
			},
		},
	}
	worker := &fakeTidbWorkerApi{}
	s := newServerForTest(pdcli, worker, nil, nil, otherKsWorker)

	rec := httptest.NewRecorder()
	req := httptest.NewRequest(http.MethodPost,
		"/api/keyspace/migrate_meta_service_group?keyspace-name=ks-a&meta-service-group-id=g1", nil)
	s.migrateMetaServiceGroup(rec, req)

	if rec.Code != http.StatusOK {
		t.Fatalf("expected 200 (other-keyspace pods ignored), got %d body=%s",
			rec.Code, rec.Body.String())
	}
}

func TestMigrateMetaServiceGroup_BlockedByConcurrentMigration(t *testing.T) {
	pdcli := &fakePDClient{
		keyspaces: map[string]*pd.KeyspaceMeta{
			"ks-a": {Name: "ks-a"},
		},
		metaServiceGroups: &[]pd.MetaServiceGroup{{ID: "g1"}},
	}
	locker := keyspacemigration.New()
	// Simulate another concurrent migration already in progress.
	release, ok := locker.TryAcquire("ks-a")
	if !ok {
		t.Fatal("setup: locker should be free")
	}
	defer release()

	worker := &fakeTidbWorkerApi{}
	s := newServerForTest(pdcli, worker, locker, nil)

	rec := httptest.NewRecorder()
	req := httptest.NewRequest(http.MethodPost,
		"/api/keyspace/migrate_meta_service_group?keyspace-name=ks-a&meta-service-group-id=g1", nil)
	s.migrateMetaServiceGroup(rec, req)

	if rec.Code != http.StatusConflict {
		t.Fatalf("expected 409, got %d body=%s", rec.Code, rec.Body.String())
	}
	if pdcli.updateKeyspaceCalled != 0 {
		t.Fatalf("must not update pd while another migration holds the lock")
	}
	if worker.pauseCalled != 0 {
		t.Fatalf("must not pause tidb-worker while another migration holds the lock")
	}
}

func TestMigrateMetaServiceGroup_RecheckCatchesRaceAndResumesWorker(t *testing.T) {
	// A TiDB shows up between the pre-check and the post-pause re-check.
	// The endpoint must roll the pause back and refuse the migration.
	pdcli := &fakePDClient{
		keyspaces: map[string]*pd.KeyspaceMeta{
			"ks-a": {Name: "ks-a"},
		},
		metaServiceGroups: &[]pd.MetaServiceGroup{{ID: "g1"}},
	}
	tm := &testTidbManager{active: map[string]testWalkInfo{}}
	tManagers := map[string]tidbmanager.TidbManager{"tidb-pool": tm}

	worker := &fakeTidbWorkerApi{}
	s := newServerForTest(pdcli, worker, nil, tManagers)
	s.tidbWorkerApi = &racingTidbWorkerApi{
		fake: worker,
		onPause: func() {
			tm.active["ks-a"] = testWalkInfo{endpoints: []string{"pod-late,10.0.0.9"}, replicas: 1}
		},
	}

	rec := httptest.NewRecorder()
	req := httptest.NewRequest(http.MethodPost,
		"/api/keyspace/migrate_meta_service_group?keyspace-name=ks-a&meta-service-group-id=g1", nil)
	s.migrateMetaServiceGroup(rec, req)

	if rec.Code != http.StatusConflict {
		t.Fatalf("expected 409 due to race, got %d body=%s", rec.Code, rec.Body.String())
	}
	if pdcli.updateKeyspaceCalled != 0 {
		t.Fatalf("must not update pd when recheck fails")
	}
	if worker.pauseCalled != 1 {
		t.Fatalf("should have paused once, got %d", worker.pauseCalled)
	}
	if worker.resumeCalled != 1 {
		t.Fatalf("should have resumed once after recheck failure, got %d", worker.resumeCalled)
	}
	if worker.pausedAt > worker.resumedAt {
		t.Fatalf("resume should have happened after pause: pausedAt=%d resumedAt=%d",
			worker.pausedAt, worker.resumedAt)
	}
}

// racingTidbWorkerApi wraps fakeTidbWorkerApi so a side effect can run after
// Pause completes, simulating a race where state changes between the pre-check
// and the post-pause recheck.
type racingTidbWorkerApi struct {
	fake    *fakeTidbWorkerApi
	onPause func()
}

func (r *racingTidbWorkerApi) Clean(ctx context.Context, keyspace string) error {
	return r.fake.Clean(ctx, keyspace)
}
func (r *racingTidbWorkerApi) Pause(ctx context.Context, keyspace string) error {
	if err := r.fake.Pause(ctx, keyspace); err != nil {
		return err
	}
	if r.onPause != nil {
		r.onPause()
	}
	return nil
}
func (r *racingTidbWorkerApi) Resume(ctx context.Context, keyspace string) error {
	return r.fake.Resume(ctx, keyspace)
}

func TestMigrateMetaServiceGroup_HappyPath(t *testing.T) {
	pdcli := &fakePDClient{
		keyspaces: map[string]*pd.KeyspaceMeta{
			"ks-a": {Name: "ks-a", Config: pd.KeyspaceConfig{MetaServiceGroupID: "g0"}},
		},
		metaServiceGroups: &[]pd.MetaServiceGroup{{ID: "g0"}, {ID: "g1"}},
	}
	worker := &fakeTidbWorkerApi{}
	locker := keyspacemigration.New()
	s := newServerForTest(pdcli, worker, locker, nil)

	rec := httptest.NewRecorder()
	req := httptest.NewRequest(http.MethodPost,
		"/api/keyspace/migrate_meta_service_group?keyspace-name=ks-a&meta-service-group-id=g1", nil)
	s.migrateMetaServiceGroup(rec, req)

	if rec.Code != http.StatusOK {
		t.Fatalf("expected 200, got %d body=%s", rec.Code, rec.Body.String())
	}
	if pdcli.updateKeyspaceCalled != 1 {
		t.Fatalf("expected exactly one pd config update, got %d", pdcli.updateKeyspaceCalled)
	}
	if pdcli.lastUpdateConfig == nil || pdcli.lastUpdateConfig.MetaServiceGroupID != "g1" {
		t.Fatalf("pd config should set meta_service_group_id to g1, got %#v", pdcli.lastUpdateConfig)
	}
	if worker.pauseCalled != 1 {
		t.Fatalf("expected exactly one pause, got %d", worker.pauseCalled)
	}
	if worker.resumeCalled != 1 {
		t.Fatalf("expected exactly one resume, got %d", worker.resumeCalled)
	}
	if worker.pausedAt >= worker.resumedAt {
		t.Fatalf("resume should follow pause: pausedAt=%d resumedAt=%d",
			worker.pausedAt, worker.resumedAt)
	}
	if locker.IsLocked("ks-a") {
		t.Fatal("lock must be released after a successful migration")
	}

	var got migrateResult
	if err := json.Unmarshal(rec.Body.Bytes(), &got); err != nil {
		t.Fatalf("decode response: %v", err)
	}
	if got.Status != "migrated" {
		t.Fatalf("want status=migrated, got %q", got.Status)
	}
	if got.MetaServiceGroupID != "g1" || got.KeyspaceName != "ks-a" {
		t.Fatalf("unexpected response: %+v", got)
	}
}

func TestMigrateMetaServiceGroup_PauseFailure(t *testing.T) {
	pdcli := &fakePDClient{
		keyspaces: map[string]*pd.KeyspaceMeta{
			"ks-a": {Name: "ks-a"},
		},
		metaServiceGroups: &[]pd.MetaServiceGroup{{ID: "g1"}},
	}
	worker := &fakeTidbWorkerApi{pauseErr: errors.New("scaler unreachable")}
	locker := keyspacemigration.New()
	s := newServerForTest(pdcli, worker, locker, nil)

	rec := httptest.NewRecorder()
	req := httptest.NewRequest(http.MethodPost,
		"/api/keyspace/migrate_meta_service_group?keyspace-name=ks-a&meta-service-group-id=g1", nil)
	s.migrateMetaServiceGroup(rec, req)

	if rec.Code != http.StatusInternalServerError {
		t.Fatalf("expected 500, got %d body=%s", rec.Code, rec.Body.String())
	}
	if pdcli.updateKeyspaceCalled != 0 {
		t.Fatalf("must not update pd if pause failed")
	}
	if locker.IsLocked("ks-a") {
		t.Fatal("lock must be released even after pause failure")
	}
}

func TestMigrateMetaServiceGroup_PdUpdateFailureResumesWorker(t *testing.T) {
	pdcli := &fakePDClient{
		keyspaces: map[string]*pd.KeyspaceMeta{
			"ks-a": {Name: "ks-a"},
		},
		metaServiceGroups: &[]pd.MetaServiceGroup{{ID: "g1"}},
		updateKeyspaceErr: errors.New("pd write rejected"),
	}
	worker := &fakeTidbWorkerApi{}
	locker := keyspacemigration.New()
	s := newServerForTest(pdcli, worker, locker, nil)

	rec := httptest.NewRecorder()
	req := httptest.NewRequest(http.MethodPost,
		"/api/keyspace/migrate_meta_service_group?keyspace-name=ks-a&meta-service-group-id=g1", nil)
	s.migrateMetaServiceGroup(rec, req)

	if rec.Code != http.StatusInternalServerError {
		t.Fatalf("expected 500, got %d body=%s", rec.Code, rec.Body.String())
	}
	if worker.pauseCalled != 1 || worker.resumeCalled != 1 {
		t.Fatalf("pd failure must resume tidb-worker; pause=%d resume=%d",
			worker.pauseCalled, worker.resumeCalled)
	}
	if locker.IsLocked("ks-a") {
		t.Fatal("lock must be released after pd failure")
	}
}