package premium_dm

import (
	"context"
	"errors"
	"fmt"
	"net/http"
	"net/http/httptest"
	"testing"
	"time"

	dm "github.com/pingcap/tiflow/dm/openapi"
	"github.com/stretchr/testify/mock"
	"github.com/stretchr/testify/require"
	center_mock "github.com/tidbcloud/dataflow-service-ng/app/event_bus/publisher/center/mock"
	modelpremiumdm "github.com/tidbcloud/dataflow-service-ng/app/models/premium_dm"
	"github.com/tidbcloud/dataflow-service-ng/app/repo"
	"github.com/tidbcloud/dataflow-service-ng/app/repo/entity"
	repomock "github.com/tidbcloud/dataflow-service-ng/app/repo/mock"
	"github.com/tidbcloud/dataflow-service-ng/idl/pbgen/proto/premium_dm/v1beta2"
	"github.com/tidbcloud/dataflow-service-ng/pkg/account/accountdb"
	"github.com/tidbcloud/dataflow-service-ng/pkg/account/accounttypes"
	"github.com/tidbcloud/dataflow-service-ng/pkg/agent"
	"github.com/tidbcloud/dataflow-service-ng/pkg/crypto"
	encryptormock "github.com/tidbcloud/dataflow-service-ng/pkg/encryptor-wrapper/mock"
	grpclimock "github.com/tidbcloud/dataflow-service-ng/pkg/grpc/grpcli/mock"
	agentpb "github.com/tidbcloud/dataflow-service/dataflow-agent-ng/idl/pbgen/proto/ng_agent"
	"github.com/tidbcloud/eventbus-datastruct/event"
	"github.com/tidbcloud/pkgv2/auth/auth0"
	"github.com/tidbcloud/pkgv2/bizerr"
	"github.com/tidbcloud/pkgv2/db"
	mock_db "github.com/tidbcloud/pkgv2/db/mock_db"
	commonv1 "github.com/tidbcloud/tidb-management-service/api/spec/global/common/v1"
	nextgenv1 "github.com/tidbcloud/tidb-management-service/api/spec/global/nextgen/v1"
	mgmt "github.com/tidbcloud/tidb-management-service/api/spec/tidb_mgmt_service/v1"
	"go.uber.org/mock/gomock"
	"google.golang.org/protobuf/encoding/protojson"
	"google.golang.org/protobuf/proto"
	"gorm.io/gorm"
)

func TestCreatePremiumMigrationRegistersClusterDeleteLock(t *testing.T) {
	ctx := premiumDMCreateTestContext()
	testDB, err := repo.GetFakeDB()
	require.NoError(t, err)

	dao := modelpremiumdm.NewPremiumDMTaskDAO(db.NewProvider(testDB))
	mgmtCli := grpclimock.NewMgmtSvcGrpcClient(t)
	encryptor := encryptormock.NewEncryptor(t)
	bizSinkEndpointRepo := repomock.NewBizSinkEndpointRepo(t)
	publisher := center_mock.NewEventCenterPublisher(t)
	service := NewPremiumDMService(
		dao,
		nil,
		mgmtCli,
		encryptor,
		agent.ClientConfig{},
		db.NewProvider(testDB),
		nil,
		bizSinkEndpointRepo,
		publisher,
	)

	req := buildCreatePremiumMigrationReq()
	cluster := buildPremiumDMCluster()
	listCall := mgmtCli.On("ListClusters", mock.Anything, mock.Anything).Return(&mgmt.ListClustersResponse{
		Clusters: []*mgmt.Cluster{cluster},
	}, nil)
	defer listCall.Unset()
	repoListCall := bizSinkEndpointRepo.On("List", mock.Anything, mock.Anything).
		Return([]entity.BizSinkEndpoint{}, nil)
	defer repoListCall.Unset()

	encryptor.On("EncryptBase64", "target-password").Return("encrypted-target-password", nil).Once()
	encryptor.On("EncryptBase64", "source-password").Return("encrypted-source-password", nil).Once()

	var registeredStatusName string
	var setStatusReq *mgmt.SetClusterCustomStatusRequest
	setStatusCall := mgmtCli.On("SetClusterCustomStatus", mock.Anything, mock.Anything).
		Run(func(args mock.Arguments) {
			setStatusReq = args.Get(1).(*mgmt.SetClusterCustomStatusRequest)
			registeredStatusName = setStatusReq.GetCustomStatus().GetStatusName()
		}).
		Return(&mgmt.ClusterCustomStatus{StatusId: 9527}, nil)
	defer setStatusCall.Unset()
	publisher.On("PublishAsync", mock.MatchedBy(func(ev *event.DisplayEvent) bool {
		return ev.DisplayEventType == string(entity.CreateDataMigration) &&
			ev.DisplayStatus == event.EventStatusStarted &&
			ev.AuthorID == 40004
	})).Once()

	resp, err := service.CreatePremiumMigration(ctx, req)
	require.NoError(t, err)
	require.NotNil(t, resp)
	require.NotNil(t, setStatusReq)
	require.NotNil(t, setStatusReq.CustomStatus)
	require.NotNil(t, setStatusReq.CustomStatus.Mutex)
	require.NotEmpty(t, registeredStatusName)
	require.Equal(t, cluster.GetClusterId(), setStatusReq.CustomStatus.GetClusterId())
	require.Equal(t, mgmt.CustomStatusType_CUSTOM_STATUS_TYPE_DATA_MIGRATION, setStatusReq.CustomStatus.GetStatusType())
	require.Equal(t, []mgmt.ClusterAction{mgmt.ClusterAction_CLUSTER_ACTION_DELETE_CLUSTER},
		setStatusReq.CustomStatus.Mutex.GetClusterAction())
	require.Equal(t, registeredStatusName, resp.GetMigrationId())

	task, bizErr := dao.GetTask(ctx, req.GetTidbId(), resp.GetMigrationId(), nil)
	require.NoError(t, bizErr)
	require.Equal(t, uint64(9527), task.ClusterStatusID)
	require.Equal(t, resp.GetMigrationId(), task.ID)
	require.Equal(t, registeredStatusName, task.ID)
	require.Equal(t, "30003", task.ProjectID)
	require.Equal(t, uint64(40004), task.CreatorID)
	require.Equal(t, "premium-dm-user@example.com", task.Creator)

	cleanupPremiumDMTask(t, testDB, task.ID)
}

func TestCreatePremiumMigrationDeletesClusterStatusOnTransactionFailure(t *testing.T) {
	ctx := premiumDMCreateTestContext()
	cluster := buildPremiumDMCluster()
	req := buildCreatePremiumMigrationReq()
	mgmtCli := grpclimock.NewMgmtSvcGrpcClient(t)
	encryptor := encryptormock.NewEncryptor(t)
	bizSinkEndpointRepo := repomock.NewBizSinkEndpointRepo(t)
	ctrl := gomock.NewController(t)
	transProvider := mock_db.NewMockTransProvider(ctrl)
	service := NewPremiumDMService(
		nil,
		nil,
		mgmtCli,
		encryptor,
		agent.ClientConfig{},
		transProvider,
		nil,
		bizSinkEndpointRepo,
		nil,
	)

	listCall := mgmtCli.On("ListClusters", mock.Anything, mock.Anything).Return(&mgmt.ListClustersResponse{
		Clusters: []*mgmt.Cluster{cluster},
	}, nil)
	defer listCall.Unset()
	repoListCall := bizSinkEndpointRepo.On("List", mock.Anything, mock.Anything).
		Return([]entity.BizSinkEndpoint{}, nil)
	defer repoListCall.Unset()
	encryptor.On("EncryptBase64", "target-password").Return("encrypted-target-password", nil).Once()
	encryptor.On("EncryptBase64", "source-password").Return("encrypted-source-password", nil).Once()
	setStatusCall := mgmtCli.On("SetClusterCustomStatus", mock.Anything, mock.Anything).
		Return(&mgmt.ClusterCustomStatus{StatusId: 8080}, nil)
	defer setStatusCall.Unset()
	var deleteStatusReq *mgmt.DeleteClusterCustomStatusRequest
	deleteStatusCall := mgmtCli.On("DeleteClusterCustomStatus", mock.Anything, mock.Anything).
		Run(func(args mock.Arguments) {
			deleteStatusReq = args.Get(1).(*mgmt.DeleteClusterCustomStatusRequest)
		}).
		Return(nil)
	defer deleteStatusCall.Unset()

	txErr := errors.New("tx failed")
	transProvider.EXPECT().Transaction(gomock.Any(), gomock.Any()).Return(txErr)

	resp, err := service.CreatePremiumMigration(ctx, req)
	require.Nil(t, resp)
	require.ErrorIs(t, err, txErr)
	require.NotNil(t, deleteStatusReq)
	require.Equal(t, cluster.GetClusterId(), deleteStatusReq.GetClusterId())
	require.Equal(t, uint64(8080), deleteStatusReq.GetStatusId())
}

func TestCreatePremiumMigrationAbortsWhenRegisteringClusterStatusFails(t *testing.T) {
	ctx := premiumDMCreateTestContext()
	cluster := buildPremiumDMCluster()
	req := buildCreatePremiumMigrationReq()
	mgmtCli := grpclimock.NewMgmtSvcGrpcClient(t)
	encryptor := encryptormock.NewEncryptor(t)
	bizSinkEndpointRepo := repomock.NewBizSinkEndpointRepo(t)
	ctrl := gomock.NewController(t)
	transProvider := mock_db.NewMockTransProvider(ctrl)
	service := NewPremiumDMService(
		nil,
		nil,
		mgmtCli,
		encryptor,
		agent.ClientConfig{},
		transProvider,
		nil,
		bizSinkEndpointRepo,
		nil,
	)

	listCall := mgmtCli.On("ListClusters", mock.Anything, mock.Anything).Return(&mgmt.ListClustersResponse{
		Clusters: []*mgmt.Cluster{cluster},
	}, nil)
	defer listCall.Unset()
	repoListCall := bizSinkEndpointRepo.On("List", mock.Anything, mock.Anything).
		Return([]entity.BizSinkEndpoint{}, nil)
	defer repoListCall.Unset()
	encryptor.On("EncryptBase64", "target-password").Return("encrypted-target-password", nil).Once()
	encryptor.On("EncryptBase64", "source-password").Return("encrypted-source-password", nil).Once()
	setStatusErr := bizerr.New(bizerr.ErrCodeDB, "register failed")
	setStatusCall := mgmtCli.On("SetClusterCustomStatus", mock.Anything, mock.Anything).
		Return((*mgmt.ClusterCustomStatus)(nil), setStatusErr)
	defer setStatusCall.Unset()

	resp, err := service.CreatePremiumMigration(ctx, req)
	require.Nil(t, resp)
	require.ErrorIs(t, err, setStatusErr)
}

func TestCreatePremiumMigrationKeepsClusterStatusOnResponseBuildFailure(t *testing.T) {
	ctx := premiumDMCreateTestContext()
	testDB, err := repo.GetFakeDB()
	require.NoError(t, err)

	const callbackName = "test:inject-invalid-premium-dm-subtask-status"
	require.NoError(t, testDB.Callback().Create().After("gorm:create").Register(callbackName, func(tx *gorm.DB) {
		task, ok := tx.Statement.Dest.(*entity.PremiumDMTask)
		if !ok || len(task.DMSources) == 0 {
			return
		}
		task.DMSources[0].SubTaskStatus = &entity.PremiumSubTaskStatus{
			LoadStatus: &dm.LoadStatus{Progress: "not-a-progress"},
		}
	}))
	defer func() {
		require.NoError(t, testDB.Callback().Create().Remove(callbackName))
	}()

	dao := modelpremiumdm.NewPremiumDMTaskDAO(db.NewProvider(testDB))
	mgmtCli := grpclimock.NewMgmtSvcGrpcClient(t)
	encryptor := encryptormock.NewEncryptor(t)
	bizSinkEndpointRepo := repomock.NewBizSinkEndpointRepo(t)
	service := NewPremiumDMService(
		dao,
		nil,
		mgmtCli,
		encryptor,
		agent.ClientConfig{},
		db.NewProvider(testDB),
		nil,
		bizSinkEndpointRepo,
		nil,
	)

	req := buildCreatePremiumMigrationReq()
	cluster := buildPremiumDMCluster()
	listCall := mgmtCli.On("ListClusters", mock.Anything, mock.Anything).Return(&mgmt.ListClustersResponse{
		Clusters: []*mgmt.Cluster{cluster},
	}, nil)
	defer listCall.Unset()
	repoListCall := bizSinkEndpointRepo.On("List", mock.Anything, mock.Anything).
		Return([]entity.BizSinkEndpoint{}, nil)
	defer repoListCall.Unset()
	encryptor.On("EncryptBase64", "target-password").Return("encrypted-target-password", nil).Once()
	encryptor.On("EncryptBase64", "source-password").Return("encrypted-source-password", nil).Once()
	setStatusCall := mgmtCli.On("SetClusterCustomStatus", mock.Anything, mock.Anything).
		Return(&mgmt.ClusterCustomStatus{StatusId: 6060}, nil)
	defer setStatusCall.Unset()

	deleteStatusCall := mgmtCli.On("DeleteClusterCustomStatus", mock.Anything, mock.Anything).Return(nil)
	defer deleteStatusCall.Unset()

	resp, err := service.CreatePremiumMigration(ctx, req)
	require.Nil(t, resp)
	require.ErrorContains(t, err, "invalid syntax")
	mgmtCli.AssertNotCalled(t, "DeleteClusterCustomStatus", mock.Anything, mock.Anything)

	var task entity.PremiumDMTask
	require.NoError(t, testDB.Where("cluster_status_id = ?", 6060).First(&task).Error)
	require.NotEmpty(t, task.ID)
	require.Equal(t, uint64(6060), task.ClusterStatusID)

	cleanupPremiumDMTask(t, testDB, task.ID)
}

func buildCreatePremiumMigrationReq() *v1beta2.CreatePremiumMigrationReq {
	now := time.Now().UnixNano()
	return &v1beta2.CreatePremiumMigrationReq{
		TidbId:      "10001",
		DisplayName: fmt.Sprintf("pdm-%d", now),
		Mode:        v1beta2.PremiumDMTaskMode_ALL,
		ImportMode:  v1beta2.PremiumDMTaskImportMode_IMPORTMODE_LOGICAL,
		Target: &v1beta2.PremiumDMTarget{
			User:     "root",
			Password: "target-password",
		},
		Sources: []*v1beta2.PremiumDMSource{{
			SourceType: v1beta2.PremiumDMSourceType_MYSQL,
			ConnProfile: &v1beta2.PremiumDMConnProfile{
				ConnType: v1beta2.PremiumDMConnType_PUBLIC,
				Endpoint: &v1beta2.PremiumDMConnProfile_Host{Host: "8.8.8.8"},
				Port:     3306,
				User:     "source-user",
				Password: "source-password",
			},
		}},
	}
}

func premiumDMCreateTestContext() context.Context {
	user := &accountdb.User{
		User: accounttypes.User{
			Updatable: accounttypes.Updatable{
				Base: accounttypes.Base{ID: 40004},
			},
			Email: "premium-dm-user@example.com",
		},
		TenantID: 20002,
	}
	return context.WithValue(context.Background(), auth0.AuthTokenClaimKey, user)
}

func buildPremiumDMCluster() *mgmt.Cluster {
	return &mgmt.Cluster{
		ClusterId:      10001,
		OrganizationId: 20002,
		ProjectId:      30003,
		CloudProvider:  string(entity.CloudProviderTypeAWS),
		Region:         "us-west-2",
		ClusterInfo: &mgmt.Cluster_NextGenCluster{
			NextGenCluster: &nextgenv1.Cluster{ServicePlan: commonv1.ServicePlan_Essential_V2},
		},
	}
}

func cleanupPremiumDMTask(t *testing.T, testDB *gorm.DB, taskID string) {
	t.Helper()
	require.NoError(t, testDB.Exec("DELETE FROM premium_dm_sources WHERE dm_task_id = ?", taskID).Error)
	require.NoError(t, testDB.Exec("DELETE FROM premium_dm_tasks WHERE id = ?", taskID).Error)
}

func TestCreatePremiumMigrationRejectsPrivatePublicSourceHost(t *testing.T) {
	req := buildCreatePremiumMigrationReq()
	req.Sources[0].ConnProfile.Endpoint = &v1beta2.PremiumDMConnProfile_Host{Host: "10.0.0.1"}
	mgmtCli := grpclimock.NewMgmtSvcGrpcClient(t)
	service := NewPremiumDMService(nil, nil, mgmtCli, nil, agent.ClientConfig{}, nil, nil, nil, nil)
	mgmtCli.On("ListClusters", mock.Anything, mock.Anything).Return(&mgmt.ListClustersResponse{
		Clusters: []*mgmt.Cluster{buildPremiumDMCluster()},
	}, nil).Once()

	resp, err := service.CreatePremiumMigration(context.Background(), req)
	require.Nil(t, resp)
	require.Error(t, err)
	require.Contains(t, err.Error(), "non-public ip")
}

func TestCreatePremiumMigrationPrecheckRejectsPrivatePublicSourceHost(t *testing.T) {
	req := buildCreatePremiumMigrationReq()
	req.Sources[0].ConnProfile.Endpoint = &v1beta2.PremiumDMConnProfile_Host{Host: "10.0.0.1"}
	mgmtCli := grpclimock.NewMgmtSvcGrpcClient(t)
	service := NewPremiumDMService(nil, nil, mgmtCli, nil, agent.ClientConfig{}, nil, nil, nil, nil)
	mgmtCli.On("ListClusters", mock.Anything, mock.Anything).Return(&mgmt.ListClustersResponse{
		Clusters: []*mgmt.Cluster{buildPremiumDMCluster()},
	}, nil).Once()

	resp, err := service.CreatePremiumMigrationPrecheck(context.Background(), req)
	require.Nil(t, resp)
	require.Error(t, err)
	require.Contains(t, err.Error(), "non-public ip")
}

func newPremiumDMServiceWithEssentialV2Cluster(t *testing.T) *PremiumDMService {
	t.Helper()
	mgmtCli := grpclimock.NewMgmtSvcGrpcClient(t)
	mgmtCli.On("ListClusters", mock.Anything, mock.Anything).Return(&mgmt.ListClustersResponse{
		Clusters: []*mgmt.Cluster{buildPremiumDMCluster()},
	}, nil)
	return NewPremiumDMService(nil, nil, mgmtCli, nil, agent.ClientConfig{}, nil, nil, nil, nil)
}

func TestValidatePremiumMigrationSourceConnRejectsPrivatePublicSourceHost(t *testing.T) {
	service := newPremiumDMServiceWithEssentialV2Cluster(t)
	req := buildPremiumMigrationSourceDBConnProfile("127.0.0.1")

	resp, err := service.ValidatePremiumMigrationSourceConn(context.Background(), req)
	require.Nil(t, resp)
	require.Error(t, err)
	require.Contains(t, err.Error(), "non-public ip")
}

func TestValidatePremiumMigrationSourceConnRejectsMissingConnProfile(t *testing.T) {
	service := newPremiumDMServiceWithEssentialV2Cluster(t)
	req := buildPremiumMigrationSourceDBConnProfile("8.8.8.8")
	req.ConnProfile = nil

	resp, err := service.ValidatePremiumMigrationSourceConn(context.Background(), req)
	require.Nil(t, resp)
	require.Error(t, err)
	require.Contains(t, err.Error(), "conn profile is required")
}

func TestValidatePremiumMigrationSourceConnPassesThroughPayloadErrorMsg(t *testing.T) {
	service, cleanup := newPremiumDMServiceWithMockAgent(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch r.URL.Path {
		case "/proxy/nextgen/resourcePools/42/dataplane/proxy/sharedPools/42/dataflow-agent/api/v1/crypto/pem":
			respBytes := marshalProtoForTest(t, &agentpb.GetPublicKeyResponse{Pem: newServicePublicKeyForTest(t)})
			_, _ = w.Write(respBytes)
		case "/proxy/nextgen/resourcePools/42/dataplane/proxy/sharedPools/42/dataflow-agent/api/v1/dm/source/exec_sql":
			respBytes := marshalProtoForTest(t, &agentpb.ExecSQLOnDMSourceResponse{
				ErrorMsg: "dial tcp 10.0.74.125:3306: i/o timeout",
			})
			_, _ = w.Write(respBytes)
		default:
			w.WriteHeader(http.StatusNotFound)
		}
	}))
	defer cleanup()

	resp, err := service.ValidatePremiumMigrationSourceConn(context.Background(), buildPremiumMigrationSourceDBConnProfile("8.8.8.8"))
	require.Nil(t, resp)
	require.Error(t, err)
	require.Equal(t, "dial tcp 10.0.74.125:3306: i/o timeout", err.Error())
	require.NotContains(t, err.Error(), "All attempts fail")
}

func TestValidatePremiumMigrationSourceConnPassesThroughGatewaySourceDBTimeout(t *testing.T) {
	service, cleanup := newPremiumDMServiceWithMockAgent(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch r.URL.Path {
		case "/proxy/nextgen/resourcePools/42/dataplane/proxy/sharedPools/42/dataflow-agent/api/v1/crypto/pem":
			respBytes := marshalProtoForTest(t, &agentpb.GetPublicKeyResponse{Pem: newServicePublicKeyForTest(t)})
			_, _ = w.Write(respBytes)
		case "/proxy/nextgen/resourcePools/42/dataplane/proxy/sharedPools/42/dataflow-agent/api/v1/dm/source/exec_sql":
			w.WriteHeader(http.StatusGatewayTimeout)
			_, _ = w.Write([]byte(`{"code":4,"message":"dial tcp 10.0.74.125:3306: i/o timeout","details":[]}`))
		default:
			w.WriteHeader(http.StatusNotFound)
		}
	}))
	defer cleanup()

	resp, err := service.ValidatePremiumMigrationSourceConn(context.Background(), buildPremiumMigrationSourceDBConnProfile("8.8.8.8"))
	require.Nil(t, resp)
	require.Error(t, err)
	require.Equal(t, "dial tcp 10.0.74.125:3306: i/o timeout", err.Error())
}

func TestValidatePremiumMigrationSourceConnKeepsTransportErrorGeneric(t *testing.T) {
	service, cleanup := newPremiumDMServiceWithMockAgent(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch r.URL.Path {
		case "/proxy/nextgen/resourcePools/42/dataplane/proxy/sharedPools/42/dataflow-agent/api/v1/crypto/pem":
			respBytes := marshalProtoForTest(t, &agentpb.GetPublicKeyResponse{Pem: newServicePublicKeyForTest(t)})
			_, _ = w.Write(respBytes)
		case "/proxy/nextgen/resourcePools/42/dataplane/proxy/sharedPools/42/dataflow-agent/api/v1/dm/source/exec_sql":
			w.WriteHeader(http.StatusGatewayTimeout)
			_, _ = w.Write([]byte(`{"code":4,"message":"upstream request timeout","details":[]}`))
		default:
			w.WriteHeader(http.StatusNotFound)
		}
	}))
	defer cleanup()

	resp, err := service.ValidatePremiumMigrationSourceConn(context.Background(), buildPremiumMigrationSourceDBConnProfile("8.8.8.8"))
	require.Nil(t, resp)
	require.Error(t, err)
	require.Equal(t, "failed to connect database", err.Error())
}

func TestGetPremiumMigrationSourceTablesRejectsPrivatePublicSourceHost(t *testing.T) {
	service := newPremiumDMServiceWithEssentialV2Cluster(t)
	req := buildPremiumMigrationSourceDBConnProfile("127.0.0.1")

	resp, err := service.GetPremiumMigrationSourceTables(context.Background(), req)
	require.Nil(t, resp)
	require.Error(t, err)
	require.Contains(t, err.Error(), "non-public ip")
}

func buildPremiumMigrationSourceDBConnProfile(host string) *v1beta2.PremiumMigrationSourceDBConnProfile {
	return &v1beta2.PremiumMigrationSourceDBConnProfile{
		TidbId: "10001",
		ConnProfile: &v1beta2.PremiumDMConnProfile{
			ConnType: v1beta2.PremiumDMConnType_PUBLIC,
			Endpoint: &v1beta2.PremiumDMConnProfile_Host{
				Host: host,
			},
			Port:     3306,
			User:     "source-user",
			Password: "source-password",
		},
	}
}

func newPremiumDMServiceWithMockAgent(t *testing.T, handler http.Handler) (*PremiumDMService, func()) {
	t.Helper()

	server := httptest.NewServer(handler)
	mgmtCli := grpclimock.NewMgmtSvcGrpcClient(t)
	mgmtCli.On("ListClusters", mock.Anything, mock.Anything).Return(&mgmt.ListClustersResponse{
		Clusters: []*mgmt.Cluster{{
			ClusterId: 10001,
			ClusterInfo: &mgmt.Cluster_NextGenCluster{
				NextGenCluster: &nextgenv1.Cluster{
					SharedPoolId: "42",
					ServicePlan:  commonv1.ServicePlan_Essential_V2,
				},
			},
		}},
	}, nil)
	mgmtCli.On("GetResourcePoolByID", mock.Anything, uint64(42)).
		Return(&nextgenv1.ResourcePool{}, nil)

	service := NewPremiumDMService(nil, nil, mgmtCli, nil, agent.ClientConfig{
		MgmtSvcProxyAddr: server.URL,
	}, nil, nil, nil, nil)

	return service, server.Close
}

func newServicePublicKeyForTest(t *testing.T) []byte {
	t.Helper()

	encryptor, err := crypto.NewRSAEncryptor(2048)
	require.NoError(t, err)
	pemBytes, err := encryptor.ExportPublicKeyAsPEM()
	require.NoError(t, err)
	return pemBytes
}

func marshalProtoForTest(t *testing.T, msg proto.Message) []byte {
	t.Helper()

	respBytes, err := protojson.Marshal(msg)
	require.NoError(t, err)
	return respBytes
}