package cmek_test

import (
	"testing"

	regionaldedicatedv1 "github.com/tidbcloud/infra-provider/api/tidbcloud/regional-server/v1"
	"github.com/tidbcloud/pkgv2/db/mock_db"
	"go.uber.org/mock/gomock"

	globalcommonv1 "github.com/tidbcloud/tidb-management-service/api/spec/global/common/v1"
	globaldedicatedv1 "github.com/tidbcloud/tidb-management-service/api/spec/global/dedicated/v1"
	"github.com/tidbcloud/tidb-management-service/dedicated-global-service/domain/adapter/account"
	"github.com/tidbcloud/tidb-management-service/dedicated-global-service/domain/adapter/account/mock_account"
	"github.com/tidbcloud/tidb-management-service/dedicated-global-service/domain/adapter/regionserver/mock_regionserver"
	"github.com/tidbcloud/tidb-management-service/dedicated-global-service/domain/entity"
	"github.com/tidbcloud/tidb-management-service/dedicated-global-service/domain/repository/cluster/mock_cluster"
	cmekrepo "github.com/tidbcloud/tidb-management-service/dedicated-global-service/domain/repository/cmek"
	"github.com/tidbcloud/tidb-management-service/dedicated-global-service/domain/repository/cmek/mock_cmek"
	"github.com/tidbcloud/tidb-management-service/dedicated-global-service/domain/repository/region/mock_region"
	"github.com/tidbcloud/tidb-management-service/dedicated-global-service/domain/service/cmek"
	"github.com/tidbcloud/tidb-management-service/utilities/bizerr"
	"github.com/tidbcloud/tidb-management-service/utilities/featuregate"
	"github.com/tidbcloud/tidb-management-service/utilities/featuregate/mock_featuregate"
	"github.com/tidbcloud/tidb-management-service/utilities/idgenerator/mock_idgenerator"
)

// The following mock packages are assumed to be generated and available.
// import "github.com/tidbcloud/pkgv2/db/mock_db"
// import "github.com/tidbcloud/infra-provider/api/tidbcloud/regional-server/v1"

func TestNewService(t *testing.T) {
	// Simple test to ensure the service constructor returns a non-nil object.
	s := cmek.NewService(nil, nil, nil, nil, nil, nil, nil, nil)
	if s == nil {
		t.Error("NewService should not return nil")
	}
}

//nolint:gocritic // ignore "too many return values"
func setupTest(t *testing.T) (
	*gomock.Controller,
	*mock_cmek.MockRepo,
	*mock_cluster.MockRepo,
	*mock_region.MockRepo,
	*mock_regionserver.MockClient,
	*mock_idgenerator.MockIDGenerator,
	*mock_account.MockClient,
	*mock_featuregate.MockService,
	cmek.CustomerManagedEncryptionKeyService,
) {
	t.Helper()

	ctrl := gomock.NewController(t)
	cmekRepo := mock_cmek.NewMockRepo(ctrl)
	clusterRepo := mock_cluster.NewMockRepo(ctrl)
	regionRepo := mock_region.NewMockRepo(ctrl)
	regionServerCli := mock_regionserver.NewMockClient(ctrl)
	idGeneratorAdapter := mock_idgenerator.NewMockIDGenerator(ctrl)
	accountCli := mock_account.NewMockClient(ctrl)
	featGateCli := mock_featuregate.NewMockService(ctrl)

	transaction := mock_db.NewFakeTransProvider()

	service := cmek.NewService(
		cmekRepo,
		clusterRepo,
		regionRepo,
		regionServerCli,
		idGeneratorAdapter,
		accountCli,
		featGateCli,
		transaction,
	)
	return ctrl, cmekRepo, clusterRepo, regionRepo, regionServerCli, idGeneratorAdapter, accountCli, featGateCli, service
}

func TestCreate(t *testing.T) {
	ctrl, cmekRepo, _, regionRepo, regionServerCli, idGeneratorAdapter, accountCli, featGateCli, service := setupTest(t)
	defer ctrl.Finish()

	ctx := t.Context()

	testAwsCmek := &entity.CustomerManagedEncryptionKey{
		OrganizationId:  1,
		ProjectId:       10,
		CloudProviderId: globalcommonv1.CloudProvider_AWS,
		RegionId:        globalcommonv1.Region_AWS_US_WEST_2,
		MasterKey: &globaldedicatedv1.CustomerManagedEncryptionKey_AwsKeyArn{
			AwsKeyArn: "arn:aws:kms:us-west-2:123456789012:key/12345678-1234-1234-1234-123456789012",
		},
	}
	testAwsReq := &globaldedicatedv1.CreateCustomerManagedEncryptionKeyRequest{
		CustomerManagedEncryptionKey: testAwsCmek,
	}

	testAzureCmek := &entity.CustomerManagedEncryptionKey{
		OrganizationId:  1,
		ProjectId:       10,
		CloudProviderId: globalcommonv1.CloudProvider_AZURE,
		RegionId:        globalcommonv1.Region_AZURE_EASTUS,
		MasterKey: &globaldedicatedv1.CustomerManagedEncryptionKey_AzureKeyVault{
			AzureKeyVault: &globaldedicatedv1.AzureKeyVault{
				AzureAppRegistration: &globaldedicatedv1.AzureAppRegistration{
					AppRegistrationId:   "app-id",
					AppRegistrationUmi:  "app-umi",
					AppRegistrationName: "app-name",
				},
				KeyVaultName: "vault-name",
				KeyVaultKey:  "key-name-valid-suffix-16s",
			},
		},
	}
	testAzureReq := &globaldedicatedv1.CreateCustomerManagedEncryptionKeyRequest{
		CustomerManagedEncryptionKey: testAzureCmek,
	}

	// Rename existing AWS tests with "aws_" prefix
	t.Run("aws_Success", func(t *testing.T) {
		accountCli.EXPECT().
			GetProject(gomock.Any(), testAwsCmek.ProjectId).
			Return(&account.ProjectInfo{AwsCmekEnabled: true}, nil).Times(1)
		featGateCli.EXPECT().CheckFeatureEnabled(
			featuregate.AwsCmek,
			testAwsCmek.OrganizationId,
			testAwsCmek.ProjectId,
			testAwsCmek.GetCloudProviderId().String(),
			testAwsCmek.RegionId.String(),
		).Return(true).Times(1)
		cmekRepo.EXPECT().Query(gomock.Any(), gomock.Any()).Return(nil, nil).Times(1)
		idGeneratorAdapter.EXPECT().MustNew().Return(uint64(123)).Times(1)
		regionRepo.EXPECT().
			SelectRegionById(gomock.Any(), testAwsCmek.RegionId, testAwsCmek.OrganizationId).
			Return(&entity.Region{
				CloudProviderId: globalcommonv1.CloudProvider_AWS,
				State:           globalcommonv1.Region_ACTIVE,
			}, nil).Times(1)
		regionRepo.EXPECT().
			GetCloudProviderById(gomock.Any(), testAwsCmek.CloudProviderId, testAwsCmek.OrganizationId).
			Return(&entity.CloudProvider{
				State: globalcommonv1.CloudProvider_ACTIVE,
			}, nil).Times(1)
		cmekRepo.EXPECT().Save(gomock.Any(), gomock.Any()).Return(nil).Times(1)

		res, berr := service.Create(ctx, testAwsReq)
		if berr != nil {
			t.Fatalf("unexpected error: %v", berr)
		}
		if res.GetId() != 123 {
			t.Errorf("expected ID 123, got %d", res.Id)
		}
	})

	t.Run("azure_Success", func(t *testing.T) {
		accountCli.EXPECT().
			GetProject(gomock.Any(), testAzureCmek.ProjectId).
			Return(&account.ProjectInfo{AwsCmekEnabled: true}, nil).Times(1)
		featGateCli.EXPECT().CheckFeatureEnabled(
			featuregate.AwsCmek,
			testAzureCmek.OrganizationId,
			testAzureCmek.ProjectId,
			testAzureCmek.GetCloudProviderId().String(),
			testAzureCmek.RegionId.String(),
		).Return(true).Times(1)
		cmekRepo.EXPECT().Query(gomock.Any(), gomock.Any()).Return(nil, nil).Times(1)
		idGeneratorAdapter.EXPECT().MustNew().Return(uint64(123)).Times(1)
		regionRepo.EXPECT().
			SelectRegionById(gomock.Any(), testAzureCmek.RegionId, testAzureCmek.OrganizationId).
			Return(&entity.Region{
				CloudProviderId: globalcommonv1.CloudProvider_AZURE,
				State:           globalcommonv1.Region_ACTIVE,
			}, nil).Times(3)
		regionRepo.EXPECT().
			GetCloudProviderById(gomock.Any(), testAzureCmek.CloudProviderId, testAzureCmek.OrganizationId).
			Return(&entity.CloudProvider{
				State: globalcommonv1.CloudProvider_ACTIVE,
			}, nil).Times(1)
		cmekRepo.EXPECT().
			GetProjectCmekKeySuffix(gomock.Any(), testAzureCmek.ProjectId).
			Return(&entity.ProjectCmekKeySuffix{
				ProjectId: testAzureCmek.ProjectId,
				Suffix:    "valid-suffix-16s",
			}, nil).
			Times(1)
		regionServerCli.EXPECT().GetAccount(gomock.Any(), gomock.Any(), gomock.Any()).
			Return(&regionaldedicatedv1.Account{
				AzureAppRegistration: &regionaldedicatedv1.AzureAppRegistration{
					AppRegistrationUmi: "app-umi",
				},
			}, nil).Times(1)
		regionServerCli.EXPECT().
			ValidateSecurityConfiguration(gomock.Any(), gomock.Any(), gomock.Any()).
			Return(&regionaldedicatedv1.ValidateSecurityConfigurationResponse{IsValid: true}, nil).Times(1)
		cmekRepo.EXPECT().Save(gomock.Any(), gomock.Any()).Return(nil).Times(1)

		res, berr := service.Create(ctx, testAzureReq)
		if berr != nil {
			t.Fatalf("unexpected error: %v", berr)
		}
		if res.GetId() != 123 {
			t.Errorf("expected ID 123, got %d", res.Id)
		}
	})

	t.Run("CmekInvalid", func(t *testing.T) {
		req := &globaldedicatedv1.CreateCustomerManagedEncryptionKeyRequest{CustomerManagedEncryptionKey: nil}
		_, berr := service.Create(ctx, req)
		if berr == nil || berr.ErrCode() != bizerr.ErrCodeParams {
			t.Errorf("expected ErrCodeParams, got %v", berr)
		}
	})

	t.Run("aws_GetProjectError", func(t *testing.T) {
		accountCli.EXPECT().
			GetProject(gomock.Any(), testAwsCmek.ProjectId).
			Return(nil, bizerr.New(bizerr.ErrCodeDB, "db error")).Times(1)
		_, berr := service.Create(ctx, testAwsReq)
		if berr == nil || berr.ErrCode() != bizerr.ErrCodeDB {
			t.Errorf("expected ErrCodeDB, got %v", berr)
		}
	})

	t.Run("azure_GetProjectError", func(t *testing.T) {
		accountCli.EXPECT().
			GetProject(gomock.Any(), testAzureCmek.ProjectId).
			Return(nil, bizerr.New(bizerr.ErrCodeDB, "db error")).Times(1)
		_, berr := service.Create(ctx, testAzureReq)
		if berr == nil || berr.ErrCode() != bizerr.ErrCodeDB {
			t.Errorf("expected ErrCodeDB, got %v", berr)
		}
	})

	t.Run("aws_CmekNotEnabled", func(t *testing.T) {
		accountCli.EXPECT().
			GetProject(gomock.Any(), testAwsCmek.ProjectId).
			Return(&account.ProjectInfo{AwsCmekEnabled: false}, nil).Times(1)
		_, berr := service.Create(ctx, testAwsReq)
		if berr == nil || berr.ErrCode() != bizerr.ErrCodeParams {
			t.Errorf("expected ErrCodeParams, got %v", berr)
		}
	})

	t.Run("azure_CmekNotEnabled", func(t *testing.T) {
		accountCli.EXPECT().
			GetProject(gomock.Any(), testAzureCmek.ProjectId).
			Return(&account.ProjectInfo{AwsCmekEnabled: false}, nil).Times(1)
		_, berr := service.Create(ctx, testAzureReq)
		if berr == nil || berr.ErrCode() != bizerr.ErrCodeParams {
			t.Errorf("expected ErrCodeParams, got %v", berr)
		}
	})

	t.Run("aws_CmekAlreadyExists", func(t *testing.T) {
		accountCli.EXPECT().
			GetProject(gomock.Any(), testAwsCmek.ProjectId).
			Return(&account.ProjectInfo{AwsCmekEnabled: true}, nil).Times(1)
		featGateCli.EXPECT().
			CheckFeatureEnabled(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).
			Return(true).Times(1)
		cmekRepo.EXPECT().
			Query(gomock.Any(), gomock.Any()).
			Return([]*entity.CustomerManagedEncryptionKey{
				{
					ProjectId: testAwsCmek.ProjectId,
					RegionId:  testAwsCmek.RegionId,
				},
			}, nil).
			Times(1)
		_, berr := service.Create(ctx, testAwsReq)
		if berr == nil || berr.ErrCode() != bizerr.ErrCodeRecordExist {
			t.Errorf("expected ErrCodeRecordExist, got %v", berr)
		}
	})

	t.Run("aws_CmekAlreadyExists_NotTheSameRegion", func(t *testing.T) {
		accountCli.EXPECT().
			GetProject(gomock.Any(), testAwsCmek.ProjectId).
			Return(&account.ProjectInfo{AwsCmekEnabled: true}, nil).Times(1)
		featGateCli.EXPECT().CheckFeatureEnabled(
			featuregate.AwsCmek,
			testAwsCmek.OrganizationId,
			testAwsCmek.ProjectId,
			testAwsCmek.GetCloudProviderId().String(),
			testAwsCmek.RegionId.String(),
		).Return(true).Times(1)
		cmekRepo.EXPECT().
			Query(gomock.Any(), gomock.Any()).
			Return([]*entity.CustomerManagedEncryptionKey{
				{
					ProjectId: testAwsCmek.ProjectId,
					RegionId:  globalcommonv1.Region_AWS_US_EAST_1,
				},
			}, nil).
			Times(1)
		idGeneratorAdapter.EXPECT().MustNew().Return(uint64(123)).Times(1)
		regionRepo.EXPECT().
			SelectRegionById(gomock.Any(), testAwsCmek.RegionId, testAwsCmek.OrganizationId).
			Return(&entity.Region{
				CloudProviderId: globalcommonv1.CloudProvider_AWS,
				State:           globalcommonv1.Region_ACTIVE,
			}, nil).Times(1)
		regionRepo.EXPECT().
			GetCloudProviderById(gomock.Any(), testAwsCmek.CloudProviderId, testAwsCmek.OrganizationId).
			Return(&entity.CloudProvider{
				State: globalcommonv1.CloudProvider_ACTIVE,
			}, nil).Times(1)
		cmekRepo.EXPECT().Save(gomock.Any(), gomock.Any()).Return(nil).Times(1)

		res, berr := service.Create(ctx, testAwsReq)
		if berr != nil {
			t.Fatalf("unexpected error: %v", berr)
		}
		if res.GetId() != 123 {
			t.Errorf("expected ID 123, got %d", res.Id)
		}
	})

	t.Run("azure_CmekAlreadyExists", func(t *testing.T) {
		accountCli.EXPECT().
			GetProject(gomock.Any(), testAzureCmek.ProjectId).
			Return(&account.ProjectInfo{AwsCmekEnabled: true}, nil).
			Times(1)
		featGateCli.EXPECT().
			CheckFeatureEnabled(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).
			Return(true).
			Times(1)
		cmekRepo.EXPECT().
			Query(gomock.Any(), gomock.Any()).
			Return([]*entity.CustomerManagedEncryptionKey{{
				ProjectId: testAzureCmek.ProjectId,
				RegionId:  testAzureCmek.RegionId,
			}}, nil).
			Times(1)
		regionRepo.EXPECT().
			SelectRegionById(gomock.Any(), testAzureCmek.RegionId, testAzureCmek.OrganizationId).
			Return(&entity.Region{
				CloudProviderId: globalcommonv1.CloudProvider_AZURE,
				State:           globalcommonv1.Region_ACTIVE,
			}, nil).Times(1)
		regionServerCli.EXPECT().GetAccount(gomock.Any(), gomock.Any(), gomock.Any()).
			Return(&regionaldedicatedv1.Account{
				AzureAppRegistration: &regionaldedicatedv1.AzureAppRegistration{
					AppRegistrationUmi: "app-umi",
				},
			}, nil).Times(1)
		_, berr := service.Create(ctx, testAzureReq)
		if berr == nil || berr.ErrCode() != bizerr.ErrCodeRecordExist {
			t.Errorf("expected ErrCodeRecordExist, got %v", berr)
		}
	})

	t.Run("aws_CheckCloudProviderAndRegionError", func(t *testing.T) {
		accountCli.EXPECT().
			GetProject(gomock.Any(), testAwsCmek.ProjectId).
			Return(&account.ProjectInfo{AwsCmekEnabled: true}, nil).Times(1)
		featGateCli.EXPECT().
			CheckFeatureEnabled(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).
			Return(true).Times(1)
		cmekRepo.EXPECT().Query(gomock.Any(), gomock.Any()).Return(nil, nil).Times(1)
		idGeneratorAdapter.EXPECT().MustNew().Return(uint64(123)).Times(1)
		regionRepo.EXPECT().
			SelectRegionById(gomock.Any(), testAwsCmek.RegionId, testAwsCmek.OrganizationId).
			Return(nil, bizerr.New(bizerr.ErrCodeRecordNotFound, "region not found")).Times(1)

		_, berr := service.Create(ctx, testAwsReq)
		if berr == nil || berr.ErrCode() != bizerr.ErrCodeRecordNotFound {
			t.Errorf("expected ErrCodeParams, got %v", berr)
		}
	})

	t.Run("azure_CheckCloudProviderAndRegionError", func(t *testing.T) {
		accountCli.EXPECT().
			GetProject(gomock.Any(), testAzureCmek.ProjectId).
			Return(&account.ProjectInfo{AwsCmekEnabled: true}, nil).Times(1)
		featGateCli.EXPECT().
			CheckFeatureEnabled(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).
			Return(true).Times(1)
		cmekRepo.EXPECT().Query(gomock.Any(), gomock.Any()).Return(nil, nil).Times(1)
		idGeneratorAdapter.EXPECT().MustNew().Return(uint64(123)).Times(1)
		regionRepo.EXPECT().
			SelectRegionById(gomock.Any(), testAzureCmek.RegionId, testAzureCmek.OrganizationId).
			Return(&entity.Region{
				CloudProviderId: globalcommonv1.CloudProvider_AZURE,
				State:           globalcommonv1.Region_ACTIVE,
			}, nil).Times(1)
		regionServerCli.EXPECT().GetAccount(gomock.Any(), gomock.Any(), gomock.Any()).
			Return(&regionaldedicatedv1.Account{
				AzureAppRegistration: &regionaldedicatedv1.AzureAppRegistration{
					AppRegistrationUmi: "app-umi",
				},
			}, nil).Times(1)
		regionRepo.EXPECT().
			SelectRegionById(gomock.Any(), testAzureCmek.RegionId, testAzureCmek.OrganizationId).
			Return(nil, bizerr.New(bizerr.ErrCodeRecordNotFound, "region not found")).Times(1)

		_, berr := service.Create(ctx, testAzureReq)
		if berr == nil || berr.ErrCode() != bizerr.ErrCodeRecordNotFound {
			t.Errorf("expected ErrCodeParams, got %v", berr)
		}
	})

	t.Run("aws_CheckCmekResourceError", func(t *testing.T) {
		testCmekInvalid := &entity.CustomerManagedEncryptionKey{
			OrganizationId:  1,
			ProjectId:       10,
			CloudProviderId: globalcommonv1.CloudProvider_AWS,
			RegionId:        globalcommonv1.Region_AWS_US_WEST_2,
			MasterKey: &globaldedicatedv1.CustomerManagedEncryptionKey_AwsKeyArn{
				AwsKeyArn: "invalid arn",
			},
		}
		accountCli.EXPECT().
			GetProject(gomock.Any(), testCmekInvalid.ProjectId).
			Return(&account.ProjectInfo{AwsCmekEnabled: true}, nil).Times(1)
		featGateCli.EXPECT().
			CheckFeatureEnabled(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).
			Return(true).Times(1)
		cmekRepo.EXPECT().Query(gomock.Any(), gomock.Any()).Return(nil, nil).Times(1)
		idGeneratorAdapter.EXPECT().MustNew().Return(uint64(123)).Times(1)
		regionRepo.EXPECT().
			SelectRegionById(gomock.Any(), testCmekInvalid.RegionId, testCmekInvalid.OrganizationId).
			Return(&entity.Region{
				CloudProviderId: globalcommonv1.CloudProvider_AWS,
				State:           globalcommonv1.Region_ACTIVE,
			}, nil).Times(1)
		regionRepo.EXPECT().
			GetCloudProviderById(gomock.Any(), testCmekInvalid.CloudProviderId, testCmekInvalid.OrganizationId).
			Return(&entity.CloudProvider{
				State: globalcommonv1.CloudProvider_ACTIVE,
			}, nil).Times(1)

		req := &globaldedicatedv1.CreateCustomerManagedEncryptionKeyRequest{
			CustomerManagedEncryptionKey: testCmekInvalid,
		}
		_, berr := service.Create(ctx, req)
		if berr == nil || berr.ErrCode() != bizerr.ErrCodeParams {
			t.Errorf("expected ErrCodeParams, got %v", berr)
		}
	})

	t.Run("azure_CheckCmekResourceError", func(t *testing.T) {
		testCmekInvalid := &entity.CustomerManagedEncryptionKey{
			OrganizationId:  1,
			ProjectId:       10,
			CloudProviderId: globalcommonv1.CloudProvider_AZURE,
			RegionId:        globalcommonv1.Region_AZURE_EASTUS,
			MasterKey: &globaldedicatedv1.CustomerManagedEncryptionKey_AzureKeyVault{
				AzureKeyVault: &globaldedicatedv1.AzureKeyVault{
					AzureAppRegistration: &globaldedicatedv1.AzureAppRegistration{
						AppRegistrationId:   "app-id",
						AppRegistrationUmi:  "app-umi",
						AppRegistrationName: "app-name",
					},
					KeyVaultName: "vault-name",
					KeyVaultKey:  "key-name-valid-suffix-16s",
				},
			},
		}
		accountCli.EXPECT().
			GetProject(gomock.Any(), testCmekInvalid.ProjectId).
			Return(&account.ProjectInfo{AwsCmekEnabled: true}, nil).Times(1)
		featGateCli.EXPECT().
			CheckFeatureEnabled(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).
			Return(true).Times(1)
		cmekRepo.EXPECT().Query(gomock.Any(), gomock.Any()).Return(nil, nil).Times(1)
		idGeneratorAdapter.EXPECT().MustNew().Return(uint64(123)).Times(1)
		regionRepo.EXPECT().
			SelectRegionById(gomock.Any(), testCmekInvalid.RegionId, testCmekInvalid.OrganizationId).
			Return(&entity.Region{
				CloudProviderId: globalcommonv1.CloudProvider_AZURE,
				State:           globalcommonv1.Region_ACTIVE,
			}, nil).Times(3)
		regionServerCli.EXPECT().GetAccount(gomock.Any(), gomock.Any(), gomock.Any()).
			Return(&regionaldedicatedv1.Account{
				AzureAppRegistration: &regionaldedicatedv1.AzureAppRegistration{
					AppRegistrationUmi: "app-umi",
				},
			}, nil).Times(1)
		regionRepo.EXPECT().
			GetCloudProviderById(gomock.Any(), testCmekInvalid.CloudProviderId, testCmekInvalid.OrganizationId).
			Return(&entity.CloudProvider{
				State: globalcommonv1.CloudProvider_ACTIVE,
			}, nil).Times(1)
		cmekRepo.EXPECT().
			GetProjectCmekKeySuffix(gomock.Any(), testCmekInvalid.ProjectId).
			Return(&entity.ProjectCmekKeySuffix{
				ProjectId: testAzureCmek.ProjectId,
				Suffix:    "valid-suffix-16s",
			}, nil).
			Times(1)
		regionServerCli.EXPECT().
			ValidateSecurityConfiguration(gomock.Any(), gomock.Any(), gomock.Any()).
			Return(&regionaldedicatedv1.ValidateSecurityConfigurationResponse{
				IsValid: false,
				Message: "invalid key",
			}, nil).Times(1)

		req := &globaldedicatedv1.CreateCustomerManagedEncryptionKeyRequest{
			CustomerManagedEncryptionKey: testCmekInvalid,
		}
		_, berr := service.Create(ctx, req)
		if berr == nil || berr.ErrCode() != bizerr.ErrCodeParams {
			t.Errorf("expected ErrCodeParams, got %v", berr)
		}
	})

	t.Run("aws_SaveError", func(t *testing.T) {
		accountCli.EXPECT().
			GetProject(gomock.Any(), testAwsCmek.ProjectId).
			Return(&account.ProjectInfo{AwsCmekEnabled: true}, nil).Times(1)
		featGateCli.EXPECT().
			CheckFeatureEnabled(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).
			Return(true).Times(1)
		cmekRepo.EXPECT().Query(gomock.Any(), gomock.Any()).Return(nil, nil).Times(1)
		idGeneratorAdapter.EXPECT().MustNew().Return(uint64(123)).Times(1)
		regionRepo.EXPECT().
			SelectRegionById(gomock.Any(), testAwsCmek.RegionId, testAwsCmek.OrganizationId).
			Return(&entity.Region{
				CloudProviderId: globalcommonv1.CloudProvider_AWS,
				State:           globalcommonv1.Region_ACTIVE,
			}, nil).Times(1)
		regionRepo.EXPECT().
			GetCloudProviderById(gomock.Any(), testAwsCmek.CloudProviderId, testAwsCmek.OrganizationId).
			Return(&entity.CloudProvider{
				State: globalcommonv1.CloudProvider_ACTIVE,
			}, nil).Times(1)
		cmekRepo.EXPECT().Save(gomock.Any(), gomock.Any()).Return(bizerr.New(bizerr.ErrCodeDB, "save error")).Times(1)

		_, berr := service.Create(ctx, testAwsReq)
		if berr == nil || berr.ErrCode() != bizerr.ErrCodeDB {
			t.Errorf("expected ErrCodeDB, got %v", berr)
		}
	})

	t.Run("azure_SaveError", func(t *testing.T) {
		accountCli.EXPECT().
			GetProject(gomock.Any(), testAzureCmek.ProjectId).
			Return(&account.ProjectInfo{AwsCmekEnabled: true}, nil).Times(1)
		featGateCli.EXPECT().
			CheckFeatureEnabled(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).
			Return(true).Times(1)
		cmekRepo.EXPECT().Query(gomock.Any(), gomock.Any()).Return(nil, nil).Times(1)
		idGeneratorAdapter.EXPECT().MustNew().Return(uint64(123)).Times(1)
		regionRepo.EXPECT().
			SelectRegionById(gomock.Any(), testAzureCmek.RegionId, testAzureCmek.OrganizationId).
			Return(&entity.Region{
				CloudProviderId: globalcommonv1.CloudProvider_AZURE,
				State:           globalcommonv1.Region_ACTIVE,
			}, nil).Times(3)
		regionServerCli.EXPECT().GetAccount(gomock.Any(), gomock.Any(), gomock.Any()).
			Return(&regionaldedicatedv1.Account{
				AzureAppRegistration: &regionaldedicatedv1.AzureAppRegistration{
					AppRegistrationUmi: "app-umi",
				},
			}, nil).Times(1)
		regionRepo.EXPECT().
			GetCloudProviderById(gomock.Any(), testAzureCmek.CloudProviderId, testAzureCmek.OrganizationId).
			Return(&entity.CloudProvider{
				State: globalcommonv1.CloudProvider_ACTIVE,
			}, nil).Times(1)
		cmekRepo.EXPECT().GetProjectCmekKeySuffix(gomock.Any(), testAzureCmek.ProjectId).
			Return(&entity.ProjectCmekKeySuffix{
				ProjectId: testAzureCmek.ProjectId,
				Suffix:    "valid-suffix-16s",
			}, nil).Times(1)
		regionServerCli.EXPECT().
			ValidateSecurityConfiguration(gomock.Any(), gomock.Any(), gomock.Any()).
			Return(&regionaldedicatedv1.ValidateSecurityConfigurationResponse{IsValid: true}, nil).
			Times(1)
		cmekRepo.EXPECT().Save(gomock.Any(), gomock.Any()).Return(bizerr.New(bizerr.ErrCodeDB, "save error")).Times(1)

		_, berr := service.Create(ctx, testAzureReq)
		if berr == nil || berr.ErrCode() != bizerr.ErrCodeDB {
			t.Errorf("expected ErrCodeDB, got %v", berr)
		}
	})

	t.Run("aws_DuplicateKeyError", func(t *testing.T) {
		accountCli.EXPECT().
			GetProject(gomock.Any(), testAwsCmek.ProjectId).
			Return(&account.ProjectInfo{AwsCmekEnabled: true}, nil).Times(1)
		featGateCli.EXPECT().
			CheckFeatureEnabled(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).
			Return(true).Times(1)
		cmekRepo.EXPECT().Query(gomock.Any(), gomock.Any()).Return(nil, nil).Times(1)
		idGeneratorAdapter.EXPECT().MustNew().Return(uint64(123)).Times(1)
		regionRepo.EXPECT().
			SelectRegionById(gomock.Any(), testAwsCmek.RegionId, testAwsCmek.OrganizationId).
			Return(&entity.Region{
				CloudProviderId: globalcommonv1.CloudProvider_AWS,
				State:           globalcommonv1.Region_ACTIVE,
			}, nil).Times(1)
		regionRepo.EXPECT().
			GetCloudProviderById(gomock.Any(), testAwsCmek.CloudProviderId, testAwsCmek.OrganizationId).
			Return(&entity.CloudProvider{
				State: globalcommonv1.CloudProvider_ACTIVE,
			}, nil).Times(1)
		cmekRepo.EXPECT().
			Save(gomock.Any(), gomock.Any()).
			Return(bizerr.New(bizerr.ErrCodeRecordExist, "duplicate key error")).Times(1)

		_, berr := service.Create(ctx, testAwsReq)
		if berr == nil || berr.ErrCode() != bizerr.ErrCodeRecordExist {
			t.Errorf("expected ErrCodeDB, got %v", berr)
		}
	})

	t.Run("azure_DuplicateKeyError", func(t *testing.T) {
		accountCli.EXPECT().
			GetProject(gomock.Any(), testAzureCmek.ProjectId).
			Return(&account.ProjectInfo{AwsCmekEnabled: true}, nil).Times(1)
		featGateCli.EXPECT().
			CheckFeatureEnabled(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).
			Return(true).Times(1)
		cmekRepo.EXPECT().Query(gomock.Any(), gomock.Any()).Return(nil, nil).Times(1)
		idGeneratorAdapter.EXPECT().MustNew().Return(uint64(123)).Times(1)
		regionRepo.EXPECT().
			SelectRegionById(gomock.Any(), testAzureCmek.RegionId, testAzureCmek.OrganizationId).
			Return(&entity.Region{
				CloudProviderId: globalcommonv1.CloudProvider_AZURE,
				State:           globalcommonv1.Region_ACTIVE,
			}, nil).Times(3)
		regionRepo.EXPECT().
			GetCloudProviderById(gomock.Any(), testAzureCmek.CloudProviderId, testAzureCmek.OrganizationId).
			Return(&entity.CloudProvider{
				State: globalcommonv1.CloudProvider_ACTIVE,
			}, nil).Times(1)
		cmekRepo.EXPECT().GetProjectCmekKeySuffix(gomock.Any(), testAzureCmek.ProjectId).
			Return(&entity.ProjectCmekKeySuffix{
				ProjectId: testAzureCmek.ProjectId,
				Suffix:    "valid-suffix-16s",
			}, nil).Times(1)
		regionServerCli.EXPECT().
			ValidateSecurityConfiguration(gomock.Any(), gomock.Any(), gomock.Any()).
			Return(&regionaldedicatedv1.ValidateSecurityConfigurationResponse{IsValid: true}, nil).
			Times(1)
		regionServerCli.EXPECT().GetAccount(gomock.Any(), gomock.Any(), gomock.Any()).
			Return(&regionaldedicatedv1.Account{
				AzureAppRegistration: &regionaldedicatedv1.AzureAppRegistration{
					AppRegistrationUmi: "app-umi",
				},
			}, nil).Times(1)
		cmekRepo.EXPECT().
			Save(gomock.Any(), gomock.Any()).
			Return(bizerr.New(bizerr.ErrCodeRecordExist, "duplicate key error")).
			Times(1)

		_, berr := service.Create(ctx, testAzureReq)
		if berr == nil || berr.ErrCode() != bizerr.ErrCodeRecordExist {
			t.Errorf("expected ErrCodeRecordExist, got %v", berr)
		}
	})

	t.Run("aws_ValidateOnly", func(t *testing.T) {
		req := &globaldedicatedv1.CreateCustomerManagedEncryptionKeyRequest{
			CustomerManagedEncryptionKey: testAwsCmek,
			ValidateOnly:                 true,
		}
		accountCli.EXPECT().
			GetProject(gomock.Any(), testAwsCmek.ProjectId).
			Return(&account.ProjectInfo{AwsCmekEnabled: true}, nil).Times(1)
		featGateCli.EXPECT().
			CheckFeatureEnabled(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).
			Return(true).Times(1)
		cmekRepo.EXPECT().Query(gomock.Any(), gomock.Any()).Return(nil, nil).Times(1)
		idGeneratorAdapter.EXPECT().MustNew().Return(uint64(123)).Times(1)
		regionRepo.EXPECT().
			SelectRegionById(gomock.Any(), testAwsCmek.RegionId, testAwsCmek.OrganizationId).
			Return(&entity.Region{
				CloudProviderId: globalcommonv1.CloudProvider_AWS,
				State:           globalcommonv1.Region_ACTIVE,
			}, nil).Times(1)
		regionRepo.EXPECT().
			GetCloudProviderById(gomock.Any(), testAwsCmek.CloudProviderId, testAwsCmek.OrganizationId).
			Return(&entity.CloudProvider{
				State: globalcommonv1.CloudProvider_ACTIVE,
			}, nil).Times(1)

		_, berr := service.Create(ctx, req)
		if berr != nil {
			t.Fatalf("unexpected error: %v", berr)
		}
	})

	t.Run("azure_ValidateOnly", func(t *testing.T) {
		req := &globaldedicatedv1.CreateCustomerManagedEncryptionKeyRequest{
			CustomerManagedEncryptionKey: testAzureCmek,
			ValidateOnly:                 true,
		}
		accountCli.EXPECT().
			GetProject(gomock.Any(), testAzureCmek.ProjectId).
			Return(&account.ProjectInfo{AwsCmekEnabled: true}, nil).Times(1)
		featGateCli.EXPECT().
			CheckFeatureEnabled(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).
			Return(true).Times(1)
		cmekRepo.EXPECT().Query(gomock.Any(), gomock.Any()).Return(nil, nil).Times(1)
		idGeneratorAdapter.EXPECT().MustNew().Return(uint64(123)).Times(1)
		regionRepo.EXPECT().
			SelectRegionById(gomock.Any(), testAzureCmek.RegionId, testAzureCmek.OrganizationId).
			Return(&entity.Region{
				CloudProviderId: globalcommonv1.CloudProvider_AZURE,
				State:           globalcommonv1.Region_ACTIVE,
			}, nil).Times(3)
		regionRepo.EXPECT().
			GetCloudProviderById(gomock.Any(), testAzureCmek.CloudProviderId, testAzureCmek.OrganizationId).
			Return(&entity.CloudProvider{
				State: globalcommonv1.CloudProvider_ACTIVE,
			}, nil).Times(1)
		cmekRepo.EXPECT().GetProjectCmekKeySuffix(gomock.Any(), testAzureCmek.ProjectId).
			Return(&entity.ProjectCmekKeySuffix{
				ProjectId: testAzureCmek.ProjectId,
				Suffix:    "valid-suffix-16s",
			}, nil).Times(1)
		regionServerCli.EXPECT().GetAccount(gomock.Any(), gomock.Any(), gomock.Any()).
			Return(&regionaldedicatedv1.Account{
				AzureAppRegistration: &regionaldedicatedv1.AzureAppRegistration{
					AppRegistrationUmi: "app-umi",
				},
			}, nil).Times(1)
		regionServerCli.EXPECT().
			ValidateSecurityConfiguration(gomock.Any(), gomock.Any(), gomock.Any()).
			Return(&regionaldedicatedv1.ValidateSecurityConfigurationResponse{IsValid: true}, nil).Times(1)

		_, berr := service.Create(ctx, req)
		if berr != nil {
			t.Fatalf("unexpected error: %v", berr)
		}
	})
}

func TestDelete(t *testing.T) {
	ctrl, cmekRepo, clusterRepo, _, _, _, _, _, service := setupTest(t)
	defer ctrl.Finish()

	ctx := t.Context()
	cmekID := uint64(123)

	t.Run("Success", func(t *testing.T) {
		cmekRepo.EXPECT().Get(gomock.Any(), cmekID).
			Return(&entity.CustomerManagedEncryptionKey{
				Id: cmekID,
			}, nil).
			Times(1)
		clusterRepo.EXPECT().Query(gomock.Any(), gomock.Any()).Return(nil, nil).Times(1)
		cmekRepo.EXPECT().Delete(gomock.Any(), cmekID).Return(nil).Times(1)

		berr := service.Delete(
			ctx,
			&globaldedicatedv1.DeleteCustomerManagedEncryptionKeyRequest{CustomerManagedEncryptionKeyId: cmekID},
		)
		if berr != nil {
			t.Fatalf("unexpected error: %v", berr)
		}
	})

	t.Run("CmekInUse", func(t *testing.T) {
		cmekRepo.EXPECT().Get(gomock.Any(), cmekID).
			Return(&entity.CustomerManagedEncryptionKey{
				Id: cmekID,
			}, nil).Times(1)
		clusterRepo.EXPECT().Query(gomock.Any(), gomock.Any()).
			Return([]*entity.Cluster{
				{},
			}, nil).
			Times(1)
		berr := service.Delete(
			ctx,
			&globaldedicatedv1.DeleteCustomerManagedEncryptionKeyRequest{CustomerManagedEncryptionKeyId: cmekID},
		)
		if berr == nil || berr.ErrCode() != bizerr.ErrCodePermission {
			t.Errorf("expected ErrCodePermission, got %v", berr)
		}
	})

	t.Run("CmekNotFound", func(t *testing.T) {
		cmekRepo.EXPECT().Get(gomock.Any(), cmekID).
			Return(nil, bizerr.New(bizerr.ErrCodeRecordNotFound, "not found")).
			Times(1)
		berr := service.Delete(
			ctx,
			&globaldedicatedv1.DeleteCustomerManagedEncryptionKeyRequest{CustomerManagedEncryptionKeyId: cmekID},
		)
		if berr == nil || berr.ErrCode() != bizerr.ErrCodeRecordNotFound {
			t.Errorf("expected ErrCodeRecordNotFound, got %v", berr)
		}
	})

	t.Run("QueryClusterError", func(t *testing.T) {
		cmekRepo.EXPECT().Get(gomock.Any(), cmekID).
			Return(&entity.CustomerManagedEncryptionKey{
				Id: cmekID,
			}, nil).Times(1)
		clusterRepo.EXPECT().Query(gomock.Any(), gomock.Any()).
			Return(nil, bizerr.New(bizerr.ErrCodeDB, "db error")).Times(1)
		berr := service.Delete(
			ctx,
			&globaldedicatedv1.DeleteCustomerManagedEncryptionKeyRequest{CustomerManagedEncryptionKeyId: cmekID},
		)
		if berr == nil || berr.ErrCode() != bizerr.ErrCodeDB {
			t.Errorf("expected ErrCodeDB, got %v", berr)
		}
	})

	t.Run("DeleteError", func(t *testing.T) {
		cmekRepo.EXPECT().Get(gomock.Any(), cmekID).
			Return(&entity.CustomerManagedEncryptionKey{Id: cmekID}, nil).Times(1)
		clusterRepo.EXPECT().Query(gomock.Any(), gomock.Any()).
			Return(nil, nil).Times(1)
		cmekRepo.EXPECT().Delete(gomock.Any(), cmekID).
			Return(bizerr.New(bizerr.ErrCodeDB, "db error")).Times(1)
		berr := service.Delete(
			ctx,
			&globaldedicatedv1.DeleteCustomerManagedEncryptionKeyRequest{CustomerManagedEncryptionKeyId: cmekID},
		)
		if berr == nil || berr.ErrCode() != bizerr.ErrCodeDB {
			t.Errorf("expected ErrCodeDB, got %v", berr)
		}
	})
}

func TestGet(t *testing.T) {
	ctrl, cmekRepo, _, _, _, _, _, _, service := setupTest(t)
	defer ctrl.Finish()

	ctx := t.Context()
	cmekID := uint64(123)
	expectedCmek := &entity.CustomerManagedEncryptionKey{Id: cmekID}

	t.Run("Success", func(t *testing.T) {
		cmekRepo.EXPECT().Get(gomock.Any(), cmekID).Return(expectedCmek, nil).Times(1)
		res, berr := service.Get(
			ctx,
			&globaldedicatedv1.GetCustomerManagedEncryptionKeyRequest{CustomerManagedEncryptionKeyId: cmekID},
		)
		if berr != nil {
			t.Fatalf("unexpected error: %v", berr)
		}
		if res.Id != expectedCmek.Id {
			t.Errorf("expected cmek with ID %d, got %d", expectedCmek.Id, res.Id)
		}
	})

	t.Run("CmekNotFound", func(t *testing.T) {
		cmekRepo.EXPECT().Get(gomock.Any(), cmekID).
			Return(nil, bizerr.New(bizerr.ErrCodeRecordNotFound, "not found")).Times(1)
		_, berr := service.Get(
			ctx,
			&globaldedicatedv1.GetCustomerManagedEncryptionKeyRequest{CustomerManagedEncryptionKeyId: cmekID},
		)
		if berr == nil || berr.ErrCode() != bizerr.ErrCodeRecordNotFound {
			t.Errorf("expected ErrCodeRecordNotFound, got %v", berr)
		}
	})

	t.Run("DBError", func(t *testing.T) {
		cmekRepo.EXPECT().Get(gomock.Any(), cmekID).
			Return(nil, bizerr.New(bizerr.ErrCodeDB, "db error")).Times(1)
		_, berr := service.Get(
			ctx,
			&globaldedicatedv1.GetCustomerManagedEncryptionKeyRequest{CustomerManagedEncryptionKeyId: cmekID},
		)
		if berr == nil || berr.ErrCode() != bizerr.ErrCodeDB {
			t.Errorf("expected ErrCodeDB, got %v", berr)
		}
	})
}

func TestList(t *testing.T) {
	ctrl, cmekRepo, _, _, _, _, _, _, service := setupTest(t)
	defer ctrl.Finish()

	ctx := t.Context()
	awsReq := &globaldedicatedv1.ListCustomerManagedEncryptionKeysRequest{
		OrganizationId:   1,
		ProjectIds:       []uint64{10},
		CloudProviderIds: []globalcommonv1.CloudProvider_Name{globalcommonv1.CloudProvider_AWS},
		RegionIds:        []globalcommonv1.Region_Name{globalcommonv1.Region_AWS_US_WEST_2},
	}
	azureReq := &globaldedicatedv1.ListCustomerManagedEncryptionKeysRequest{
		OrganizationId:   1,
		ProjectIds:       []uint64{10},
		CloudProviderIds: []globalcommonv1.CloudProvider_Name{globalcommonv1.CloudProvider_AZURE},
		RegionIds:        []globalcommonv1.Region_Name{globalcommonv1.Region_AZURE_EASTUS},
	}

	expectedCmeks := []*entity.CustomerManagedEncryptionKey{{Id: 123}, {Id: 456}}

	t.Run("aws_Success", func(t *testing.T) {
		cmekRepo.EXPECT().Query(gomock.Any(), &cmekrepo.Filter{
			CloudProviderNames: awsReq.CloudProviderIds,
			RegionNames:        awsReq.RegionIds,
			ProjectIds:         awsReq.ProjectIds,
			OrganizationIds:    []uint64{awsReq.OrganizationId},
		}).Return(expectedCmeks, nil).Times(1)
		res, berr := service.List(ctx, awsReq)
		if berr != nil {
			t.Fatalf("unexpected error: %v", berr)
		}
		if len(res) != 2 {
			t.Errorf("expected 2 cmeks, got %d", len(res))
		}
	})

	t.Run("azure_Success", func(t *testing.T) {
		cmekRepo.EXPECT().Query(gomock.Any(), &cmekrepo.Filter{
			CloudProviderNames: azureReq.CloudProviderIds,
			RegionNames:        azureReq.RegionIds,
			ProjectIds:         azureReq.ProjectIds,
			OrganizationIds:    []uint64{azureReq.OrganizationId},
		}).Return(expectedCmeks, nil).Times(1)
		res, berr := service.List(ctx, azureReq)
		if berr != nil {
			t.Fatalf("unexpected error: %v", berr)
		}
		if len(res) != 2 {
			t.Errorf("expected 2 cmeks, got %d", len(res))
		}
	})

	t.Run("DBError", func(t *testing.T) {
		cmekRepo.EXPECT().Query(gomock.Any(), gomock.Any()).
			Return(nil, bizerr.New(bizerr.ErrCodeDB, "db error")).Times(1)
		_, berr := service.List(ctx, awsReq)
		if berr == nil || berr.ErrCode() != bizerr.ErrCodeDB {
			t.Errorf("expected ErrCodeDB, got %v", berr)
		}
	})
}

func TestGetCmekAccessIamPrincipal(t *testing.T) {
	ctrl, cmekRepo, _, regionRepo, regionServerCli, _, accountCli, featGateCli, service := setupTest(t)
	defer ctrl.Finish()

	ctx := t.Context()
	awsReq := &globaldedicatedv1.GetCmekAccessIamPrincipalRequest{
		OrganizationId:  1,
		ProjectId:       10,
		CloudProviderId: globalcommonv1.CloudProvider_AWS,
		RegionId:        globalcommonv1.Region_AWS_US_WEST_2,
	}

	azureReq := &globaldedicatedv1.GetCmekAccessIamPrincipalRequest{
		OrganizationId:  1,
		ProjectId:       10,
		CloudProviderId: globalcommonv1.CloudProvider_AZURE,
		RegionId:        globalcommonv1.Region_AZURE_EASTUS,
	}

	t.Run("Success_AWS_NewAccount", func(t *testing.T) {
		regionRepo.EXPECT().
			GetCloudProviderById(gomock.Any(), awsReq.CloudProviderId, awsReq.OrganizationId).
			Return(&entity.CloudProvider{
				Id:    globalcommonv1.CloudProvider_AWS,
				State: globalcommonv1.CloudProvider_ACTIVE,
			}, nil).
			Times(1)
		regionRepo.EXPECT().
			SelectRegionById(gomock.Any(), awsReq.RegionId, awsReq.OrganizationId).
			Return(&entity.Region{
				CloudProviderId: globalcommonv1.CloudProvider_AWS,
				State:           globalcommonv1.Region_ACTIVE,
			}, nil).
			Times(2)
		accountCli.EXPECT().
			GetProject(gomock.Any(), awsReq.ProjectId).
			Return(&account.ProjectInfo{AwsCmekEnabled: true}, nil).
			Times(1)
		featGateCli.EXPECT().
			CheckFeatureEnabled(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).
			Return(true).
			Times(1)
		regionServerCli.EXPECT().
			GetAccount(gomock.Any(), gomock.Any(), gomock.Any()).
			Return(nil, bizerr.New(bizerr.ErrCodeRecordNotFound, "not found")).
			Times(1)
		regionServerCli.EXPECT().
			AssignAccount(gomock.Any(), gomock.Any(), gomock.Any()).
			Return(&regionaldedicatedv1.Account{
				AccountId: "aws-account-1",
			}, nil).
			Times(1)

		principal, berr := service.GetCmekAccessIamPrincipal(ctx, awsReq)
		if berr != nil {
			t.Fatalf("unexpected error: %v", berr)
		}
		if principal.AwsPrincipalAccountId != "aws-account-1" {
			t.Errorf("expected account ID 'aws-account-1', got %s", principal.AwsPrincipalAccountId)
		}
	})

	t.Run("Success_AWS_ExistingAccount", func(t *testing.T) {
		regionRepo.EXPECT().
			GetCloudProviderById(gomock.Any(), awsReq.CloudProviderId, awsReq.OrganizationId).
			Return(&entity.CloudProvider{
				Id:    globalcommonv1.CloudProvider_AWS,
				State: globalcommonv1.CloudProvider_ACTIVE,
			}, nil).
			Times(1)
		regionRepo.EXPECT().
			SelectRegionById(gomock.Any(), awsReq.RegionId, awsReq.OrganizationId).
			Return(&entity.Region{
				CloudProviderId: globalcommonv1.CloudProvider_AWS,
				State:           globalcommonv1.Region_ACTIVE,
			}, nil).
			Times(2)
		accountCli.EXPECT().
			GetProject(gomock.Any(), awsReq.ProjectId).
			Return(&account.ProjectInfo{AwsCmekEnabled: true}, nil).
			Times(1)
		featGateCli.EXPECT().
			CheckFeatureEnabled(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).
			Return(true).
			Times(1)
		regionServerCli.EXPECT().
			GetAccount(gomock.Any(), gomock.Any(), gomock.Any()).
			Return(&regionaldedicatedv1.Account{
				AccountId: "aws-account-2",
			}, nil).
			Times(1)

		principal, berr := service.GetCmekAccessIamPrincipal(ctx, awsReq)
		if berr != nil {
			t.Fatalf("unexpected error: %v", berr)
		}
		if principal.AwsPrincipalAccountId != "aws-account-2" {
			t.Errorf("expected account ID 'aws-account-2', got %s", principal.AwsPrincipalAccountId)
		}
	})

	t.Run("Success_Azure_NewAccount", func(t *testing.T) {
		regionRepo.EXPECT().
			GetCloudProviderById(gomock.Any(), azureReq.CloudProviderId, azureReq.OrganizationId).
			Return(&entity.CloudProvider{
				Id:    globalcommonv1.CloudProvider_AZURE,
				State: globalcommonv1.CloudProvider_ACTIVE,
			}, nil).
			Times(1)
		regionRepo.EXPECT().
			SelectRegionById(gomock.Any(), azureReq.RegionId, azureReq.OrganizationId).
			Return(&entity.Region{
				CloudProviderId: globalcommonv1.CloudProvider_AZURE,
				State:           globalcommonv1.Region_ACTIVE,
			}, nil).
			Times(2)
		accountCli.EXPECT().
			GetProject(gomock.Any(), azureReq.ProjectId).
			Return(&account.ProjectInfo{AwsCmekEnabled: true}, nil).
			Times(1)
		featGateCli.EXPECT().
			CheckFeatureEnabled(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).
			Return(true).
			Times(1)
		regionServerCli.EXPECT().
			GetAccount(gomock.Any(), gomock.Any(), gomock.Any()).
			Return(nil, bizerr.New(bizerr.ErrCodeRecordNotFound, "not found")).Times(1)
		regionServerCli.EXPECT().
			AssignAccount(gomock.Any(), gomock.Any(), gomock.Any()).
			Return(&regionaldedicatedv1.Account{
				AzureAppRegistration: &regionaldedicatedv1.AzureAppRegistration{
					AppRegistrationId:   "azure-app-id",
					AppRegistrationName: "azure-app-name",
					AppRegistrationUmi:  "azure-app-umi",
				},
			}, nil).
			Times(1)
		cmekRepo.EXPECT().GetProjectCmekKeySuffix(gomock.Any(), azureReq.ProjectId).Return(&entity.ProjectCmekKeySuffix{
			Suffix: "test-suffix",
		}, nil).Times(1)

		principal, berr := service.GetCmekAccessIamPrincipal(ctx, azureReq)
		if berr != nil {
			t.Fatalf("unexpected error: %v", berr)
		}
		if principal.AzureCmekKeySuffix != "test-suffix" {
			t.Errorf("expected suffix 'test-suffix', got %s", principal.AzureCmekKeySuffix)
		}
	})

	t.Run("Success_Azure_Existing", func(t *testing.T) {
		regionRepo.EXPECT().
			GetCloudProviderById(gomock.Any(), azureReq.CloudProviderId, azureReq.OrganizationId).
			Return(&entity.CloudProvider{
				Id:    globalcommonv1.CloudProvider_AZURE,
				State: globalcommonv1.CloudProvider_ACTIVE,
			}, nil).
			Times(1)
		regionRepo.EXPECT().
			SelectRegionById(gomock.Any(), azureReq.RegionId, azureReq.OrganizationId).
			Return(&entity.Region{
				CloudProviderId: globalcommonv1.CloudProvider_AZURE,
				State:           globalcommonv1.Region_ACTIVE,
			}, nil).
			Times(2)
		accountCli.EXPECT().
			GetProject(gomock.Any(), azureReq.ProjectId).
			Return(&account.ProjectInfo{AwsCmekEnabled: true}, nil).
			Times(1)
		featGateCli.EXPECT().
			CheckFeatureEnabled(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).
			Return(true).
			Times(1)
		regionServerCli.EXPECT().
			GetAccount(gomock.Any(), gomock.Any(), gomock.Any()).
			Return(&regionaldedicatedv1.Account{
				AzureAppRegistration: &regionaldedicatedv1.AzureAppRegistration{
					AppRegistrationId:   "azure-app-id",
					AppRegistrationName: "azure-app-name",
					AppRegistrationUmi:  "azure-app-umi",
				},
			}, nil).
			Times(1)
		cmekRepo.EXPECT().GetProjectCmekKeySuffix(gomock.Any(), azureReq.ProjectId).Return(&entity.ProjectCmekKeySuffix{
			Suffix: "test-suffix",
		}, nil).Times(1)

		principal, berr := service.GetCmekAccessIamPrincipal(ctx, azureReq)
		if berr != nil {
			t.Fatalf("unexpected error: %v", berr)
		}
		if principal.AzureCmekKeySuffix != "test-suffix" {
			t.Errorf("expected suffix 'test-suffix', got %s", principal.AzureCmekKeySuffix)
		}
	})
}