package regionalservice import ( "context" "encoding/base64" "encoding/json" "fmt" "strings" "testing" "github.com/tidbcloud/observability/pkg/resourcemanager" "github.com/tidbcloud/observability/test/mockresourcemanager" "go.uber.org/mock/gomock" prometheusv1 "github.com/prometheus-operator/prometheus-operator/pkg/apis/monitoring/v1" v1 "github.com/tidbcloud/observability/api/regionalserver/resourceworker/v1" "github.com/tidbcloud/observability/pkg/common/configs" "github.com/tidbcloud/observability/pkg/common/errors" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/util/intstr" "github.com/tidbcloud/observability/pkg/common/constants" "github.com/tidbcloud/observability/pkg/common/message" "github.com/tidbcloud/observability/pkg/common/structs" "github.com/stretchr/testify/assert" ) type option func(*message.CreateClusterReq) func mockReq(opts ...option) message.CreateClusterReq { req := message.CreateClusterReq{ ClusterBaseInfo: structs.ClusterBaseInfo{ CommonReqInfo: structs.CommonReqInfo{ Vendor: constants.AWS, Region: "us-east-2", }, TenantID: "tenant111", ProjectID: "project111", ApplicationID: "application111", BizType: "tidbcloud", ID: "id111", Tags: []string{"tag111"}, Version: "v111", ShareMode: constants.EXCLUSIVE, Note: "note111", }, AgentAccessARN: "useraaaa", MonitorSpecID: "default", } for _, opt := range opts { opt(&req) } return req } func mockSharedReq(opts ...option) message.CreateClusterReq { req := message.CreateClusterReq{ ClusterBaseInfo: structs.ClusterBaseInfo{ CommonReqInfo: structs.CommonReqInfo{ Vendor: constants.AWS, Region: "us-east-2", }, TenantID: "tenant222", ProjectID: "project222", ApplicationID: "application222", BizType: "tidbcloud", ID: "id222", Tags: []string{"tag222"}, Version: "v222", ShareMode: constants.SHARED, Note: "note222", }, AgentAccessARN: "useraaaa", ParentClusterID: "parentID", KeySpaceID: "777", MonitorSpecID: "", } for _, opt := range opts { opt(&req) } return req } func withAgentAWS() option { return func(req *message.CreateClusterReq) { req.AgentEnvironment = &structs.AgentEnvironment{ AWS: &structs.AWSAgentEnvironment{ VpcID: "vpc111", NodeSubnetIDs: []string{"subnet111"}, SecurityGroupID: "sg111", Account: "account111", OIDC: "oidc.eks.us-west-2.amazonaws.com/id/TEST1", }, } } } func withAgentAliCloud() option { return func(req *message.CreateClusterReq) { req.AgentEnvironment = &structs.AgentEnvironment{ AliCloud: &structs.AliCloudAgentEnvironment{ VpcID: "vpc111", NodeSubnetIDs: []string{"subnet111"}, SecurityGroupID: "sg111", Account: "account111", OIDC: "oidc.eks.us-west-2.amazonaws.com/id/TEST1", }, } } } func withAgentGCP() option { return func(req *message.CreateClusterReq) { req.AgentEnvironment = &structs.AgentEnvironment{ GCP: &structs.GCPAgentEnvironment{ GCPProject: "project111", NetworkName: "network111", }, } } } func withObservedTiDB() option { return func(req *message.CreateClusterReq) { req.Observed = structs.Observed{ TiDBClusterTopology: &structs.ClusterTopologyType{ TiDB: structs.TiDB{ Replicas: 2, }, TiKV: structs.TiKV{ Replicas: 3, }, PD: structs.PD{ Replicas: 3, }, TiFlash: &structs.TiFlash{ Replicas: 1, }, }, } } } func withAgentAzure() option { return func(req *message.CreateClusterReq) { req.AgentEnvironment = &structs.AgentEnvironment{ Azure: &structs.AzureAgentEnvironment{ OIDC: "oidc111", SubscriptionID: "sub111", ResourceGroup: "rg111", }, } } } func TestBuildAgentRoleRequirements(t *testing.T) { req := mockReq() t.Run("normal", func(t *testing.T) { iam, err := BuildAgentRoleRequirements(req) assert.NoError(t, err) assert.Equal(t, "o11y-agent-dev-tidbcloud-tenant111-us-east-2", iam.Name) assert.NotEmpty(t, iam.Spec.Aws.AssumeRolePolicyDocument) assert.Equal(t, 1, len(iam.Spec.Aws.Policies)) }) t.Run("build GCP", func(t *testing.T) { origin := req.Vendor req.Vendor = constants.GCP req.AgentAccessARN = "test@gcp-autonomous-dev.iam.gserviceaccount.com" defer func() { req.Vendor = origin }() original := configs.GetServiceInfo() defer configs.Init(original) configs.Init(structs.ServiceInfo{ ServerName: constants.ObservabilityRegionalServer, Port: 8080, Vendor: string(constants.GCP), Region: defaultRegion, Host: "127.0.0.1", Env: constants.EnvDev, MonitorSpec: map[string]structs.MonitorSpec{ "default": { Cpu: "10m", Memory: "200Mi", InitialStorage: "40Gi", OverSell: 1, IsDefault: true, }, }, AZConfig: map[string]structs.AvailabilityZoneConfig{ "zone1": { EndpointServiceID: "e1", IngressClassName: "i1", }, "zone2": { EndpointServiceID: "e2", IngressClassName: "i2", }, }, AccountInfo: structs.AccountInfoConfig{ Oidc: "gcp-autonomous-dev.svc.id.goog", Account: "gcp-autonomous-dev.iam.gserviceaccount.com", Region: "us-west1", }, BizTypeConfigs: map[string]structs.BizTypeConfig{ "b": { AgentManagerARN: "aaa", }, "tidbcloud": { AgentManagerARN: "role2", }, }, }) iam, err := BuildAgentRoleRequirements(req) bucket := ObjectStorageResourceName(req.ClusterBaseInfo) assert.NoError(t, err) assert.Equal(t, "o11y-agent-dev-tidbcloud-tenant111-us-east-2", iam.Name) assert.Equal(t, "tidbcloud-tenant111", iam.Namespace) assert.True(t, iam.Spec.Gcp.UseExternalPrincipal) assert.Equal(t, 1, len(iam.Spec.Gcp.AssumeRoleAccounts)) assert.Equal(t, 1, len(iam.Spec.Gcp.Policies[v1.BuildGCPPolicyKey(v1.GCPResourceTypeCloudStorage, bucket)])) assert.Equal(t, req.ID, iam.Spec.Gcp.Policies[v1.BuildGCPPolicyKey(v1.GCPResourceTypeCloudStorage, bucket)][0].PolicyName) assert.Equal(t, "", iam.Spec.Gcp.Policies[v1.BuildGCPPolicyKey(v1.GCPResourceTypeCloudStorage, bucket)][0].PolicyCondition) }) t.Run("build GCP shared role with external principal", func(t *testing.T) { originVendor := req.Vendor originSharedRole := req.AgentSharedRole originAgentEnv := req.AgentEnvironment originAccessARN := req.AgentAccessARN req.Vendor = constants.GCP req.AgentSharedRole = true req.AgentEnvironment = nil req.AgentAccessARN = "test-shared@gcp-autonomous-dev.iam.gserviceaccount.com" defer func() { req.Vendor = originVendor req.AgentSharedRole = originSharedRole req.AgentEnvironment = originAgentEnv req.AgentAccessARN = originAccessARN }() original := configs.GetServiceInfo() defer configs.Init(original) configs.Init(structs.ServiceInfo{ Vendor: string(constants.GCP), Region: defaultRegion, Env: constants.EnvDev, AccountInfo: structs.AccountInfoConfig{ Oidc: "gcp-autonomous-dev.svc.id.goog", Account: "gcp-autonomous-dev.iam.gserviceaccount.com", Region: "us-west1", }, BizTypeConfigs: map[string]structs.BizTypeConfig{ "tidbcloud": { AgentManagerARN: "role2", }, }, }) iam, err := BuildAgentRoleRequirements(req) assert.NoError(t, err) assert.True(t, iam.Spec.Gcp.UseExternalPrincipal) assert.Equal(t, []string{"serviceAccount:test-shared@gcp-autonomous-dev.iam.gserviceaccount.com"}, iam.Spec.Gcp.AssumeRoleAccounts) bucket := ObjectStorageResourceName(req.ClusterBaseInfo) policies := iam.Spec.Gcp.Policies[v1.BuildGCPPolicyKey(v1.GCPResourceTypeCloudStorage, bucket)] assert.Equal(t, 1, len(policies)) assert.Equal(t, "", policies[0].PolicyCondition) }) t.Run("build GCP shared role with managed agent", func(t *testing.T) { originVendor := req.Vendor originSharedRole := req.AgentSharedRole originAgentEnv := req.AgentEnvironment originAccessARN := req.AgentAccessARN req.Vendor = constants.GCP req.AgentSharedRole = true req.AgentEnvironment = &structs.AgentEnvironment{ GCP: &structs.GCPAgentEnvironment{ GCPProject: "project111", }, } req.AgentAccessARN = "test-managed@gcp-autonomous-dev.iam.gserviceaccount.com" defer func() { req.Vendor = originVendor req.AgentSharedRole = originSharedRole req.AgentEnvironment = originAgentEnv req.AgentAccessARN = originAccessARN }() original := configs.GetServiceInfo() defer configs.Init(original) configs.Init(structs.ServiceInfo{ Vendor: string(constants.GCP), Region: defaultRegion, Env: constants.EnvDev, AccountInfo: structs.AccountInfoConfig{ Oidc: "gcp-autonomous-dev.svc.id.goog", Account: "gcp-autonomous-dev.iam.gserviceaccount.com", Region: "us-west1", }, BizTypeConfigs: map[string]structs.BizTypeConfig{ "tidbcloud": { AgentManagerARN: "role2", }, }, }) iam, err := BuildAgentRoleRequirements(req) assert.NoError(t, err) assert.False(t, iam.Spec.Gcp.UseExternalPrincipal) assert.Equal(t, 2, len(iam.Spec.Gcp.AssumeRoleAccounts)) assert.Equal(t, "serviceAccount:gcp-autonomous-dev.svc.id.goog[tidbcloud-tenant111/tidbcloud-tenant111]", iam.Spec.Gcp.AssumeRoleAccounts[0]) assert.Equal(t, "serviceAccount:test-managed@gcp-autonomous-dev.iam.gserviceaccount.com", iam.Spec.Gcp.AssumeRoleAccounts[1]) bucket := ObjectStorageResourceName(req.ClusterBaseInfo) policies := iam.Spec.Gcp.Policies[v1.BuildGCPPolicyKey(v1.GCPResourceTypeCloudStorage, bucket)] assert.Equal(t, 1, len(policies)) assert.Equal(t, "", policies[0].PolicyCondition) }) t.Run("wrong vendor", func(t *testing.T) { origin := req.Vendor req.Vendor = "UNKNOWN" defer func() { req.Vendor = origin }() original := configs.GetServiceInfo() defer configs.Init(original) configs.Init(structs.ServiceInfo{ ServerName: constants.ObservabilityRegionalServer, Port: 8080, Vendor: "UNKNOWN", Region: defaultRegion, Host: "127.0.0.1", Env: constants.EnvDev, MonitorSpec: map[string]structs.MonitorSpec{ "default": { Cpu: "10m", Memory: "200Mi", InitialStorage: "40Gi", OverSell: 1, IsDefault: true, }, }, AZConfig: map[string]structs.AvailabilityZoneConfig{ "zone1": { EndpointServiceID: "e1", IngressClassName: "i1", }, "zone2": { EndpointServiceID: "e2", IngressClassName: "i2", }, }, AccountInfo: structs.AccountInfoConfig{ Oidc: "gcp-autonomous-dev.svc.id.goog", Account: "gcp-autonomous-dev.iam.gserviceaccount.com", Region: "us-west1", }, BizTypeConfigs: map[string]structs.BizTypeConfig{ "b": { AgentManagerARN: "aaa", }, "tidbcloud": { AgentManagerARN: "role2", }, }, }) _, err := BuildAgentRoleRequirements(req) assert.Error(t, err) assert.Equal(t, errors.ParameterInvalid, err.(errors.Error).Code()) }) t.Run("build Azure", func(t *testing.T) { origin := req.Vendor req.Vendor = constants.Azure req.AgentAccessARN = "azure-node-principal-id" req.AgentEnvironment = nil defer func() { req.Vendor = origin }() original := configs.GetServiceInfo() defer configs.Init(original) configs.Init(structs.ServiceInfo{ Vendor: string(constants.Azure), Region: "eastus", Env: constants.EnvDev, AZConfig: map[string]structs.AvailabilityZoneConfig{ "zone1": { EndpointServiceID: "e1", IngressClassName: "i1", }, }, AccountInfo: structs.AccountInfoConfig{ Region: "eastus", SubscriptionID: "test", StorageAccount: "test", BlobStorageResourceGroup: "test", }, BizTypeConfigs: map[string]structs.BizTypeConfig{ "b": { AgentManagerARN: "aaa", }, "tidbcloud": { AgentManagerARN: "role2", }, }, }) iam, err := BuildAgentRoleRequirements(req) assert.NoError(t, err) assert.Equal(t, "o11y-agent-dev-tidbcloud-tenant111-eastus", iam.Name) assert.True(t, iam.Spec.Azure.UseExternalPrincipal) assert.Empty(t, iam.Spec.Azure.WorkloadServiceAccounts) assert.Equal(t, []string{"azure-node-principal-id"}, iam.Spec.Azure.AdditionalRoleMembers) assert.Empty(t, iam.Spec.Azure.RoleCondition) }) t.Run("build Azure with managed identity", func(t *testing.T) { origin := req.Vendor req.Vendor = constants.Azure req.AgentAccessARN = "azure-node-principal-id" req.AgentEnvironment = &structs.AgentEnvironment{ Azure: &structs.AzureAgentEnvironment{ OIDC: "oidc111", SubscriptionID: "sub111", ResourceGroup: "rg111", }, } defer func() { req.Vendor = origin req.AgentEnvironment = nil }() original := configs.GetServiceInfo() defer configs.Init(original) configs.Init(structs.ServiceInfo{ Vendor: string(constants.Azure), Region: "eastus", Env: constants.EnvDev, AZConfig: map[string]structs.AvailabilityZoneConfig{ "zone1": { EndpointServiceID: "e1", IngressClassName: "i1", }, }, AccountInfo: structs.AccountInfoConfig{ Region: "eastus", SubscriptionID: "test", StorageAccount: "test", BlobStorageResourceGroup: "test", }, BizTypeConfigs: map[string]structs.BizTypeConfig{ "tidbcloud": { AgentManagerARN: "role2", }, }, }) iam, err := BuildAgentRoleRequirements(req) assert.NoError(t, err) assert.False(t, iam.Spec.Azure.UseExternalPrincipal) assert.NotEmpty(t, iam.Spec.Azure.WorkloadServiceAccounts) assert.NotEmpty(t, iam.Spec.Azure.RoleCondition) }) t.Run("build AliCloud", func(t *testing.T) { origin := req.Vendor req.Vendor = constants.AliCloud req.AgentAccessARN = "test-arn" defer func() { req.Vendor = origin }() original := configs.GetServiceInfo() defer configs.Init(original) configs.Init(structs.ServiceInfo{ ServerName: constants.ObservabilityRegionalServer, Port: 8080, Vendor: string(constants.AliCloud), Region: defaultRegion, Host: "127.0.0.1", Env: constants.EnvDev, MonitorSpec: map[string]structs.MonitorSpec{ "default": { Cpu: "10m", Memory: "200Mi", InitialStorage: "40Gi", OverSell: 1, IsDefault: true, }, }, AZConfig: map[string]structs.AvailabilityZoneConfig{ "zone1": { EndpointServiceID: "e1", IngressClassName: "i1", }, "zone2": { EndpointServiceID: "e2", IngressClassName: "i2", }, }, AccountInfo: structs.AccountInfoConfig{ Oidc: "aaaa", Account: "bbbbb", Region: "ap-southeast-1", }, BizTypeConfigs: map[string]structs.BizTypeConfig{ "b": { AgentManagerARN: "aaa", }, "tidbcloud": { AgentManagerARN: "role2", }, }, }) iam, err := BuildAgentRoleRequirements(req) assert.NoError(t, err) assert.Equal(t, "o11y-agent-dev-tidbcloud-tenant111-us-east-2", iam.Name) assert.Equal(t, "tidbcloud-tenant111", iam.Namespace) assert.Equal(t, 1, len(iam.Spec.AliCloud.Policies)) assert.NotEmpty(t, iam.Spec.AliCloud.AssumeRolePolicyDocument) }) } func TestBuildClusterRoleBindingResourceRequirements(t *testing.T) { t.Run("normal", func(t *testing.T) { req := mockReq() clusterRole, err := BuildClusterRoleBindingResourceRequirements(req) assert.NoError(t, err) assert.Equal(t, "tidbcloud-tenant111", clusterRole.Name) assert.Empty(t, clusterRole.Annotations) assert.NotEmpty(t, clusterRole.Labels) assert.Equal(t, "tenant111", clusterRole.Labels[constants.LabelTenantID]) }) } func TestBuildAgentRoleRequirements_Account(t *testing.T) { req := mockReq() original := configs.GetServiceInfo() defer configs.Init(original) configs.Init(structs.ServiceInfo{ ServerName: constants.ObservabilityRegionalServer, Port: 8080, Vendor: string(constants.GCP), Region: defaultRegion, Host: "127.0.0.1", Env: constants.EnvDev, MonitorSpec: map[string]structs.MonitorSpec{ "default": { Cpu: "10m", Memory: "200Mi", InitialStorage: "40Gi", OverSell: 1, IsDefault: true, }, }, AZConfig: map[string]structs.AvailabilityZoneConfig{ "zone1": { EndpointServiceID: "e1", IngressClassName: "i1", }, "zone2": { EndpointServiceID: "e2", IngressClassName: "i2", }, }, AccountInfo: structs.AccountInfoConfig{ Oidc: "gcp-autonomous-dev.svc.id.goog", Account: "gcp-autonomous-dev.iam.gserviceaccount.com", Region: "us-west1", }, BizTypeConfigs: map[string]structs.BizTypeConfig{ "b": { AgentManagerARN: "aaa", }, "tidbcloud": { AgentManagerARN: "role2", }, }, }) origin := req.Vendor req.Vendor = constants.GCP defer func() { req.Vendor = origin }() t.Run("wrong account", func(t *testing.T) { req1 := req req1.AgentAccessARN = "test" _, err := BuildAgentRoleRequirements(req1) assert.Error(t, err) }) t.Run("wrong account 1", func(t *testing.T) { req1 := req req1.AgentAccessARN = "test@iam.gserviceaccount.com" _, err := BuildAgentRoleRequirements(req1) assert.Error(t, err) }) t.Run("service account", func(t *testing.T) { originArn := req.AgentAccessARN req.AgentAccessARN = "regional-server@gcp-autonomous-dev.iam.gserviceaccount.com" defer func() { req.AgentAccessARN = originArn }() iam, err := BuildAgentRoleRequirements(req) assert.NoError(t, err) assert.Equal(t, 1, len(iam.Spec.Gcp.AssumeRoleAccounts)) assert.True(t, strings.Contains(iam.Spec.Gcp.AssumeRoleAccounts[0], GCPSAPrefix)) assert.Equal(t, "serviceAccount:regional-server@gcp-autonomous-dev.iam.gserviceaccount.com", iam.Spec.Gcp.AssumeRoleAccounts[0]) }) t.Run("update", func(t *testing.T) { originArn := req.AgentAccessARN req.AgentAccessARN = "regional-server@gcp-autonomous-dev.iam.gserviceaccount.com" defer func() { req.AgentAccessARN = originArn }() iam, err := BuildAgentRoleRequirements(req) assert.NoError(t, err) assert.Equal(t, 1, len(iam.Spec.Gcp.AssumeRoleAccounts)) assert.True(t, strings.Contains(iam.Spec.Gcp.AssumeRoleAccounts[0], GCPSAPrefix)) assert.Equal(t, "serviceAccount:regional-server@gcp-autonomous-dev.iam.gserviceaccount.com", iam.Spec.Gcp.AssumeRoleAccounts[0]) newARN := "serviceAccount:regional-server@gcp-autonomous-dev.iam.gserviceaccount.com" newIAM, update, err := AppendAssumeEntity(context.TODO(), iam, newARN) assert.NoError(t, err) assert.False(t, update) assert.Equal(t, 1, len(newIAM.Spec.Gcp.AssumeRoleAccounts)) newARN = "regional-test@gcp-autonomous-dev.iam.gserviceaccount.com" newIAM, update, err = AppendAssumeEntity(context.TODO(), iam, newARN) assert.NoError(t, err) assert.True(t, update) assert.Equal(t, 2, len(newIAM.Spec.Gcp.AssumeRoleAccounts)) }) } func TestBuildAgentDataPlaneRoleRequirements(t *testing.T) { original := configs.GetServiceInfo() defer configs.Init(original) t.Run("aws-normal", func(t *testing.T) { configs.Init(structs.ServiceInfo{ Vendor: string(constants.AWS), Env: constants.EnvDev, Region: "us-west-2", }) req := mockSharedReq(withAgentAWS()) iam, err := BuildAgentDataPlaneRoleRequirements(req) assert.NoError(t, err) assert.Equal(t, "o11y-agent-data-dev-tidbcloud-tenant222-us-west-2-account111", iam.Name) assert.NotEmpty(t, iam.Spec.Aws.AssumeRolePolicyDocument) assert.Equal(t, 1, len(iam.Spec.Aws.Policies)) }) t.Run("gcp-normal", func(t *testing.T) { configs.Init(structs.ServiceInfo{ Vendor: string(constants.GCP), Env: constants.EnvDev, Region: "us-west1", }) req := mockSharedReq(withAgentGCP()) iam, err := BuildAgentDataPlaneRoleRequirements(req) assert.NoError(t, err) assert.Equal(t, "o11y-agent-data-dev-tidbcloud-tenant222-us-west1", iam.Name) assert.NotEmpty(t, iam.Spec.Gcp.AssumeRoleAccounts) }) t.Run("azure-normal", func(t *testing.T) { configs.Init(structs.ServiceInfo{ Vendor: string(constants.Azure), Env: constants.EnvDev, Region: "eastus", }) req := mockSharedReq(withAgentAzure()) iam, err := BuildAgentDataPlaneRoleRequirements(req) assert.NoError(t, err) assert.Equal(t, "o11y-agent-data-dev-tidbcloud-tenant222-project222-eastus", iam.Name) assert.NotEmpty(t, iam.Spec.Azure.WorkloadServiceAccounts) }) t.Run("alicloud-normal", func(t *testing.T) { configs.Init(structs.ServiceInfo{ Vendor: string(constants.AliCloud), Env: constants.EnvDev, Region: "us-west-2", }) req := mockSharedReq(withAgentAliCloud()) req.Vendor = "alicloud" iam, err := BuildAgentDataPlaneRoleRequirements(req) assert.NoError(t, err) assert.Equal(t, "o11y-agent-data-dev-tidbcloud-tenant222-us-west-2", iam.Name) assert.NotEmpty(t, iam.Spec.AliCloud.AssumeRolePolicyDocument) assert.Equal(t, 1, len(iam.Spec.AliCloud.Policies)) }) } func TestBuildNamespaceResourceRequirements(t *testing.T) { t.Run("normal", func(t *testing.T) { ob, _, err := BuildNamespaceResourceRequirements(context.TODO(), "biz1", "tenant1", "shared", false) assert.NoError(t, err) assert.Equal(t, "biz1-tenant1", ob.Name) }) t.Run("empty", func(t *testing.T) { _, _, err := BuildNamespaceResourceRequirements(context.TODO(), "", "aaa", "shared", false) assert.Error(t, err) _, _, err = BuildNamespaceResourceRequirements(context.TODO(), "aaa", "", "shared", false) assert.Error(t, err) }) t.Run("multi tenancy", func(t *testing.T) { original := configs.GetServiceInfo() defer configs.Init(original) configs.Init(structs.ServiceInfo{ BizTypeConfigs: map[string]structs.BizTypeConfig{ "b": { BizType: "b", MultiTenancyRule: true, }, }, }) ob, _, err := BuildNamespaceResourceRequirements(context.TODO(), "b", "tenant1", "shared", false) assert.NoError(t, err) assert.Equal(t, "b-tenant1", ob.Name) assert.Equal(t, "true", ob.Labels[constants.LabelMultiTenancyRule]) }) t.Run("nextgen enabled", func(t *testing.T) { ctrl := gomock.NewController(t) defer ctrl.Finish() dataAccountRW := mockresourcemanager.NewMockDataAccountReaderWriter(ctrl) resourcemanager.MockDataAccountReaderWriter(dataAccountRW) dataAccountRW.EXPECT().UniqueInt32ForCluster(gomock.Any()).Return(int32(12345), nil).Times(1) ob, vmAccountID, err := BuildNamespaceResourceRequirements(context.TODO(), "tidbcloud", "tenant123", constants.SHARED, true) assert.NoError(t, err) assert.Equal(t, "tidbcloud-tenant123", ob.Name) assert.Equal(t, "12345", vmAccountID) assert.Equal(t, "12345", ob.Labels[constants.LabelVMAccountID]) assert.Equal(t, "tidbcloud", ob.Labels[constants.LabelBizType]) assert.Equal(t, "tenant123", ob.Labels[constants.LabelTenantID]) }) t.Run("nextgen enabled with multi tenancy", func(t *testing.T) { ctrl := gomock.NewController(t) defer ctrl.Finish() dataAccountRW := mockresourcemanager.NewMockDataAccountReaderWriter(ctrl) resourcemanager.MockDataAccountReaderWriter(dataAccountRW) original := configs.GetServiceInfo() defer configs.Init(original) configs.Init(structs.ServiceInfo{ BizTypeConfigs: map[string]structs.BizTypeConfig{ "tidbcloud": { BizType: "tidbcloud", MultiTenancyRule: true, }, }, }) dataAccountRW.EXPECT().UniqueInt32ForCluster(gomock.Any()).Return(int32(67890), nil).Times(1) ob, vmAccountID, err := BuildNamespaceResourceRequirements(context.TODO(), "tidbcloud", "tenant456", constants.SHARED, true) assert.NoError(t, err) assert.Equal(t, "tidbcloud-tenant456", ob.Name) assert.Equal(t, "67890", vmAccountID) assert.Equal(t, "67890", ob.Labels[constants.LabelVMAccountID]) assert.Equal(t, "true", ob.Labels[constants.LabelMultiTenancyRule]) assert.Equal(t, "tidbcloud", ob.Labels[constants.LabelBizType]) assert.Equal(t, "tenant456", ob.Labels[constants.LabelTenantID]) }) t.Run("nextgen enabled but uniqueInt32 fails", func(t *testing.T) { ctrl := gomock.NewController(t) defer ctrl.Finish() dataAccountRW := mockresourcemanager.NewMockDataAccountReaderWriter(ctrl) resourcemanager.MockDataAccountReaderWriter(dataAccountRW) dataAccountRW.EXPECT().UniqueInt32ForCluster(gomock.Any()).Return(int32(0), fmt.Errorf("failed to generate unique ID")).Times(1) _, _, err := BuildNamespaceResourceRequirements(context.TODO(), "tidbcloud", "tenant789", constants.SHARED, true) assert.Error(t, err) assert.Contains(t, err.Error(), "failed to generate unique ID") }) } func TestBuildObjectStorageResourceRequirements(t *testing.T) { t.Run("normal", func(t *testing.T) { ob, err := BuildObjectStorageResourceRequirements("bucket111") assert.NoError(t, err) assert.Equal(t, "bucket111", ob.Name) assert.Equal(t, "bucket111", ob.Spec.Name) }) t.Run("empty", func(t *testing.T) { _, err := BuildObjectStorageResourceRequirements("") assert.Error(t, err) }) } func TestBuildClusterResourceRequirements(t *testing.T) { ctrl := gomock.NewController(t) defer ctrl.Finish() dataAccountRW := mockresourcemanager.NewMockDataAccountReaderWriter(ctrl) resourcemanager.MockDataAccountReaderWriter(dataAccountRW) clusterRW := mockresourcemanager.NewMockClusterReaderWriter(ctrl) resourcemanager.MockClusterReaderWriter(clusterRW) // Also mock RDS cluster reader writer for RDS mode tests rdsClusterRW := mockresourcemanager.NewMockClusterReaderWriter(ctrl) resourcemanager.MockRDSClusterReaderWriter(rdsClusterRW) iamRW := mockresourcemanager.NewMockIAMReaderWriter(ctrl) resourcemanager.MockIAMReaderWriter(iamRW) dataAccountRW.EXPECT().UniqueInt32ForCluster(gomock.Any()).Return(int32(111), nil).AnyTimes() t.Run("normal -- legacy role not found", func(t *testing.T) { req := mockReq(withAgentAWS()) iamRW.EXPECT().IsRoleExists(gomock.Any(), "o11y-agent-data-dev-tidbcloud-tenant111-us-east-2", gomock.Any()).Return(false, nil).Times(1) iamRW.EXPECT().GetRole(gomock.Any(), "o11y-agent-data-dev-tidbcloud-tenant111-us-east-2-account111", gomock.Any()).Return(&v1.IAM{ ObjectMeta: metav1.ObjectMeta{ Name: "o11y-agent-data-dev-tidbcloud-tenant111-us-east-2-account111", }, Spec: v1.IAMSpec{ Aws: v1.AWSIAM{ AssumeRolePolicyDocument: "test", }, }, }, nil).Times(1) iamRW.EXPECT().GetRole(gomock.Any(), "o11y-agent-dev-tidbcloud-tenant111-us-east-2", gomock.Any()).Return(&v1.IAM{ ObjectMeta: metav1.ObjectMeta{ Name: "o11y-agent-dev-tidbcloud-tenant111-us-east-2 ", }, Spec: v1.IAMSpec{ Aws: v1.AWSIAM{ AssumeRolePolicyDocument: "test", }, }, }, nil).Times(1) cluster, err := BuildClusterResourceRequirements(context.TODO(), req, "", nil) assert.NoError(t, err) assert.Equal(t, req.ID, cluster.Name) assert.NotEmpty(t, cluster.Annotations) assert.NotEmpty(t, cluster.Labels) assert.Equal(t, "tenant111", cluster.Labels[constants.LabelTenantID]) assert.Equal(t, string(constants.EXCLUSIVE), cluster.Annotations[constants.AnnotationShareMode]) assert.Equal(t, int32(2), cluster.Spec.Prometheus.Replicas) assert.Equal(t, "v2.33.5", cluster.Spec.Prometheus.Version) assert.Equal(t, "15d", string(cluster.Spec.Prometheus.Retention)) assert.NotEmpty(t, cluster.Spec.Prometheus.Query) assert.Equal(t, *cluster.Spec.Prometheus.Query.MaxConcurrency, int32(10)) assert.Empty(t, cluster.Spec.VictoriaAccount) }) t.Run("normal -- legacy role found", func(t *testing.T) { req := mockReq(withAgentAWS()) iamRW.EXPECT().IsRoleExists(gomock.Any(), "o11y-agent-data-dev-tidbcloud-tenant111-us-east-2", gomock.Any()).Return(true, nil).Times(1) iamRW.EXPECT().GetRole(gomock.Any(), "o11y-agent-data-dev-tidbcloud-tenant111-us-east-2", gomock.Any()).Return(&v1.IAM{ ObjectMeta: metav1.ObjectMeta{ Name: "o11y-agent-data-dev-tidbcloud-tenant111-us-east-2", }, Spec: v1.IAMSpec{ Aws: v1.AWSIAM{ Account: "account111", AssumeRolePolicyDocument: "test", }, }, }, nil).Times(2) iamRW.EXPECT().GetRole(gomock.Any(), "o11y-agent-dev-tidbcloud-tenant111-us-east-2", gomock.Any()).Return(&v1.IAM{ ObjectMeta: metav1.ObjectMeta{ Name: "o11y-agent-dev-tidbcloud-tenant111-us-east-2 ", }, Spec: v1.IAMSpec{ Aws: v1.AWSIAM{ AssumeRolePolicyDocument: "test", }, }, }, nil).Times(1) cluster, err := BuildClusterResourceRequirements(context.TODO(), req, "", nil) assert.NoError(t, err) assert.Equal(t, req.ID, cluster.Name) assert.NotEmpty(t, cluster.Annotations) assert.NotEmpty(t, cluster.Labels) assert.Equal(t, "tenant111", cluster.Labels[constants.LabelTenantID]) assert.Equal(t, string(constants.EXCLUSIVE), cluster.Annotations[constants.AnnotationShareMode]) assert.Equal(t, int32(2), cluster.Spec.Prometheus.Replicas) assert.Equal(t, "v2.33.5", cluster.Spec.Prometheus.Version) assert.Equal(t, "15d", string(cluster.Spec.Prometheus.Retention)) assert.NotEmpty(t, cluster.Spec.Prometheus.Query) assert.Equal(t, *cluster.Spec.Prometheus.Query.MaxConcurrency, int32(10)) assert.Empty(t, cluster.Spec.VictoriaAccount) }) t.Run("shared", func(t *testing.T) { clusterRW.EXPECT().GetCluster(gomock.Any(), "parentID").Return(&v1.Cluster{ Spec: v1.ClusterSpec{ VictoriaAccount: v1.VictoriaAccount{ AccountID: 999, ProjectID: 0, }, }, }, nil).AnyTimes() req := mockSharedReq() cluster, err := BuildClusterResourceRequirements(context.TODO(), req, "", nil) assert.NoError(t, err) assert.Equal(t, req.ID, cluster.Name) assert.NotEmpty(t, cluster.Annotations) assert.NotEmpty(t, cluster.Labels) assert.Empty(t, cluster.Spec.Prometheus.Resources) assert.Equal(t, "tenant222", cluster.Labels[constants.LabelTenantID]) assert.Equal(t, string(constants.SHARED), cluster.Annotations[constants.AnnotationShareMode]) assert.Equal(t, int32(999), cluster.Spec.VictoriaAccount.AccountID) assert.NotEmpty(t, cluster.Spec.VictoriaAccount.AvailabilityZone) assert.Empty(t, cluster.Spec.Prometheus) //assert.Equal(t, req.AgentAccessARN, cluster.Spec.Endpoint.UserARN) }) t.Run("3-layers gcp", func(t *testing.T) { originalService := configs.GetServiceInfo() defer func() { configs.Init(originalService) }() configs.Init(structs.ServiceInfo{ ServerName: constants.ObservabilityRegionalServer, Port: 8080, Vendor: string(constants.GCP), Region: defaultRegion, Host: "127.0.0.1", Env: constants.EnvDev, MonitorSpec: map[string]structs.MonitorSpec{ "default": { Cpu: "10m", Memory: "200Mi", InitialStorage: "40Gi", OverSell: 1, IsDefault: true, }, }, AZConfig: map[string]structs.AvailabilityZoneConfig{ "zone1": { EndpointServiceID: "e1", IngressClassName: "i1", DefaultDataAZ: true, }, }, AccountInfo: structs.AccountInfoConfig{ Oidc: "oidc.eks.us-west-2.amazonaws.com/id/005D92AA81DE2DCE14A2991BAC54FEF8", OidcARN: "arn:aws:iam::552185537300:oidc-provider/oidc.eks.us-west-2.amazonaws.com/id/005D92AA81DE2DCE14A2991BAC54FEF8", Account: "aaaa", Region: "us-west-2", }, VMClusterInfo: []structs.VictoriaMetricsCluster{ { VMClusterName: "test1", }, }, BizTypeConfigs: map[string]structs.BizTypeConfig{ "b": { AgentManagerARN: "aaa", }, "tidbcloud": { AgentManagerARN: "role2", }, }, }) clusterRW.EXPECT().GetCluster(gomock.Any(), "parentID").Return(&v1.Cluster{ Spec: v1.ClusterSpec{ VictoriaAccount: v1.VictoriaAccount{ AccountID: 999, ProjectID: 0, }, }, }, nil).AnyTimes() iamRW.EXPECT().GetRole(gomock.Any(), gomock.Any(), gomock.Any()).Return(&v1.IAM{ ObjectMeta: metav1.ObjectMeta{ Labels: map[string]string{ constants.LabelAlias: "test1", }, }, Spec: v1.IAMSpec{ Gcp: v1.GCPIAM{ ProjectID: "test1", }, }, }, nil).Times(2) req := mockSharedReq(withAgentGCP(), withObservedTiDB()) cluster, err := BuildClusterResourceRequirements(context.TODO(), req, "", nil) assert.NoError(t, err) assert.Equal(t, req.ID, cluster.Name) assert.NotEmpty(t, cluster.Annotations) assert.NotEmpty(t, cluster.Labels) assert.Empty(t, cluster.Spec.Prometheus.Resources) assert.Equal(t, "tenant222", cluster.Labels[constants.LabelTenantID]) assert.Equal(t, string(constants.SHARED), cluster.Annotations[constants.AnnotationShareMode]) assert.Equal(t, int32(999), cluster.Spec.VictoriaAccount.AccountID) assert.NotEmpty(t, cluster.Spec.VictoriaAccount.AvailabilityZone) assert.Empty(t, cluster.Spec.Prometheus) //assert.Equal(t, req.AgentAccessARN, cluster.Spec.Endpoint.UserARN) }) t.Run("nextgen enabled", func(t *testing.T) { ctrl := gomock.NewController(t) defer ctrl.Finish() dataAccountRW := mockresourcemanager.NewMockDataAccountReaderWriter(ctrl) resourcemanager.MockDataAccountReaderWriter(dataAccountRW) dataAccountRW.EXPECT().UniqueInt32ForCluster(gomock.Any()).Return(int32(12345), nil).Times(1) ob, vmAccountID, err := BuildNamespaceResourceRequirements(context.TODO(), "tidbcloud", "tenant123", constants.SHARED, true) assert.NoError(t, err) assert.Equal(t, "tidbcloud-tenant123", ob.Name) assert.Equal(t, "12345", vmAccountID) assert.Equal(t, "12345", ob.Labels[constants.LabelVMAccountID]) assert.Equal(t, "tidbcloud", ob.Labels[constants.LabelBizType]) assert.Equal(t, "tenant123", ob.Labels[constants.LabelTenantID]) }) t.Run("nextgen enabled with multi tenancy", func(t *testing.T) { ctrl := gomock.NewController(t) defer ctrl.Finish() dataAccountRW := mockresourcemanager.NewMockDataAccountReaderWriter(ctrl) resourcemanager.MockDataAccountReaderWriter(dataAccountRW) original := configs.GetServiceInfo() defer configs.Init(original) configs.Init(structs.ServiceInfo{ BizTypeConfigs: map[string]structs.BizTypeConfig{ "tidbcloud": { BizType: "tidbcloud", MultiTenancyRule: true, }, }, }) dataAccountRW.EXPECT().UniqueInt32ForCluster(gomock.Any()).Return(int32(67890), nil).Times(1) ob, vmAccountID, err := BuildNamespaceResourceRequirements(context.TODO(), "tidbcloud", "tenant456", constants.SHARED, true) assert.NoError(t, err) assert.Equal(t, "tidbcloud-tenant456", ob.Name) assert.Equal(t, "67890", vmAccountID) assert.Equal(t, "67890", ob.Labels[constants.LabelVMAccountID]) assert.Equal(t, "true", ob.Labels[constants.LabelMultiTenancyRule]) assert.Equal(t, "tidbcloud", ob.Labels[constants.LabelBizType]) assert.Equal(t, "tenant456", ob.Labels[constants.LabelTenantID]) }) t.Run("shared with next gen", func(t *testing.T) { dataAccountFactory := mockresourcemanager.NewMockDataAccountReaderWriterFactory(ctrl) dataAccountRW := mockresourcemanager.NewMockDataAccountReaderWriter(ctrl) dataAccountRW.EXPECT().UniqueInt32ForCluster(gomock.Any()).Return(int32(111), nil).AnyTimes() dataAccountFactory.EXPECT().GetReaderWriter(gomock.Any()).Return(dataAccountRW).AnyTimes() resourcemanager.MockDataAccountReaderWriterFactory(dataAccountFactory) req := mockSharedReq() req.ParentClusterID = "" req.Vendor = constants.AWS req.AgentAccessAccount = "123456" req.AgentNextGen = true cluster, err := BuildClusterResourceRequirements(context.TODO(), req, "41", nil) assert.NoError(t, err) assert.Equal(t, req.ID, cluster.Name) assert.NotEmpty(t, cluster.Annotations) assert.NotEmpty(t, cluster.Labels) assert.Empty(t, cluster.Spec.Prometheus.Resources) assert.Equal(t, "tenant222", cluster.Labels[constants.LabelTenantID]) assert.Equal(t, string(constants.SHARED), cluster.Annotations[constants.AnnotationShareMode]) assert.Equal(t, int32(41), cluster.Spec.VictoriaAccount.AccountID) assert.Equal(t, int32(111), cluster.Spec.VictoriaAccount.ProjectID) }) t.Run("shared with next gen azure agent access account", func(t *testing.T) { originalService := configs.GetServiceInfo() defer func() { configs.Init(originalService) }() configs.Init(structs.ServiceInfo{ Vendor: string(constants.Azure), Region: "westus3", AZConfig: map[string]structs.AvailabilityZoneConfig{ "zone1": { EndpointServiceID: "e1", IngressClassName: "i1", DefaultDataAZ: true, }, }, AccountInfo: structs.AccountInfoConfig{ Region: "westus3", SubscriptionID: "sub111", StorageAccount: "storage111", BlobStorageResourceGroup: "rg111", }, VMClusterInfo: []structs.VictoriaMetricsCluster{ { VMClusterName: "test1", }, }, }) req := mockSharedReq() req.ParentClusterID = "" req.Vendor = constants.Azure req.Region = "westus3" req.AgentAccessAccount = "ad881159-0593-4502-96e4-e2698b790a3d" req.AgentNextGen = true cluster, err := BuildClusterResourceRequirements(context.TODO(), req, "41", nil) assert.NoError(t, err) assert.Equal(t, req.AgentAccessAccount, cluster.Spec.Endpoint.UserARN) }) t.Run("mute alert should result in empty template id", func(t *testing.T) { // Set up mock for VM space (needed for shared mode) clusterRW.EXPECT().GetCluster(gomock.Any(), "parentIDskipalertrule").Return(&v1.Cluster{ Spec: v1.ClusterSpec{ VictoriaAccount: v1.VictoriaAccount{ AccountID: 999, }, }, }, nil).Times(1) // Create request with parent cluster ID (required for skipAlertRule) req := mockSharedReq() req.ParentClusterID = "parentIDskipalertrule" req.BizType = "tidbcloud" req.ShareMode = constants.SHARED req.KeySpaceID = "1" // Set KeySpaceID to 1 so ProjectID becomes 1 req.MuteAlert = true cluster, err := BuildClusterResourceRequirements(context.TODO(), req, "", nil) assert.NoError(t, err) assert.Equal(t, req.ID, cluster.Name) assert.NotEmpty(t, cluster.Annotations) assert.NotEmpty(t, cluster.Labels) assert.Empty(t, cluster.Spec.Prometheus.Resources) assert.Equal(t, string(constants.SHARED), cluster.Annotations[constants.AnnotationShareMode]) assert.NotEmpty(t, cluster.Spec.VictoriaAccount.AvailabilityZone) assert.Empty(t, cluster.Spec.Prometheus) //assert.Equal(t, req.AgentAccessARN, cluster.Spec.Endpoint.UserARN) // Verify that the LabelTemplateID is NOT present when skipAlertRule is true templateID, _ := cluster.Labels[constants.LabelTemplateID] assert.Empty(t, templateID) }) t.Run("serverless child cluster with parentCluster in RDS mode - inherits namespace", func(t *testing.T) { originalService := configs.GetServiceInfo() defer func() { configs.Init(originalService) }() configs.Init(structs.ServiceInfo{ ServerName: constants.ObservabilityRegionalServer, Port: 8080, Vendor: string(constants.AWS), Region: defaultRegion, Host: "127.0.0.1", Env: constants.EnvDev, MonitorSpec: map[string]structs.MonitorSpec{ "default": { Cpu: "10m", Memory: "200Mi", InitialStorage: "40Gi", OverSell: 1, IsDefault: true, }, }, AZConfig: map[string]structs.AvailabilityZoneConfig{ "zone1": { EndpointServiceID: "e1", IngressClassName: "i1", DefaultDataAZ: true, }, }, AccountInfo: structs.AccountInfoConfig{ Oidc: "oidc.eks.us-west-2.amazonaws.com/id/005D92AA81DE2DCE14A2991BAC54FEF8", OidcARN: "arn:aws:iam::552185537300:oidc-provider/oidc.eks.us-west-2.amazonaws.com/id/005D92AA81DE2DCE14A2991BAC54FEF8", Account: "aaaa", Region: "us-west-2", }, VMClusterInfo: []structs.VictoriaMetricsCluster{ { VMClusterName: "test1", }, }, BizTypeConfigs: map[string]structs.BizTypeConfig{ "tidbcloud": { AgentManagerARN: "role2", }, }, }) // Create a parent cluster object to pass parentCluster := &v1.Cluster{ ObjectMeta: metav1.ObjectMeta{ Namespace: "parent-namespace", }, Spec: v1.ClusterSpec{ VictoriaAccount: v1.VictoriaAccount{ AccountID: 888, ProjectID: 0, }, }, } // Set up mock for getVMSpace internal GetCluster call (uses RDS reader in RDS mode) rdsClusterRW.EXPECT().GetCluster(gomock.Any(), "parentClusterRDS").Return(parentCluster, nil).Times(1) // Create request with parent cluster ID and KeySpaceID (serverless child cluster) req := mockSharedReq() req.ParentClusterID = "parentClusterRDS" req.BizType = "tidbcloud" req.ShareMode = constants.SHARED req.KeySpaceID = "123" // Create context with RDS mode enabled ctx := context.WithValue(context.TODO(), constants.UseRDSKey, true) cluster, err := BuildClusterResourceRequirements(ctx, req, "", parentCluster) assert.NoError(t, err) assert.Equal(t, req.ID, cluster.Name) // Verify the cluster inherits the parent's namespace in RDS mode assert.Equal(t, "parent-namespace", cluster.Namespace, "Child cluster should inherit parent's namespace in RDS mode") }) t.Run("serverless child cluster without RDS mode - uses default namespace", func(t *testing.T) { parentCluster := &v1.Cluster{ ObjectMeta: metav1.ObjectMeta{ Namespace: "parent-namespace", }, Spec: v1.ClusterSpec{ VictoriaAccount: v1.VictoriaAccount{ AccountID: 777, ProjectID: 0, }, }, } // Set up mock for getVMSpace internal GetCluster call clusterRW.EXPECT().GetCluster(gomock.Any(), "parentClusterNoRDS").Return(parentCluster, nil).Times(1) // Create request with parent cluster ID and KeySpaceID req := mockSharedReq() req.ParentClusterID = "parentClusterNoRDS" req.BizType = "tidbcloud" req.ShareMode = constants.SHARED req.KeySpaceID = "456" // Create context without RDS mode (default) ctx := context.TODO() cluster, err := BuildClusterResourceRequirements(ctx, req, "", parentCluster) assert.NoError(t, err) assert.Equal(t, req.ID, cluster.Name) // Verify the cluster uses the default namespace (not parent's) when not in RDS mode assert.Equal(t, NamespaceName(req.BizType, req.TenantID), cluster.Namespace, "Child cluster should use default namespace when not in RDS mode") }) t.Run("cluster with nil parentCluster uses default namespace", func(t *testing.T) { clusterRW.EXPECT().GetCluster(gomock.Any(), "parentID").Return(&v1.Cluster{ Spec: v1.ClusterSpec{ VictoriaAccount: v1.VictoriaAccount{ AccountID: 999, ProjectID: 0, }, }, }, nil).AnyTimes() req := mockSharedReq() cluster, err := BuildClusterResourceRequirements(context.TODO(), req, "", nil) assert.NoError(t, err) assert.Equal(t, req.ID, cluster.Name) // Verify that when parentCluster is nil, default namespace is used assert.Equal(t, NamespaceName(req.BizType, req.TenantID), cluster.Namespace, "Cluster should use default namespace when parentCluster is nil") }) } func Test_pickAZ(t *testing.T) { t.Run("normal", func(t *testing.T) { azIndex = -1 k1, v1, err := pickAZ() assert.NoError(t, err) k2, v2, err := pickAZ() assert.NoError(t, err) k3, v3, err := pickAZ() assert.NoError(t, err) assert.NotEqual(t, k1, k2) assert.Equal(t, k1, k3) assert.NotEqual(t, v1.EndpointServiceID, v2.EndpointServiceID) assert.Equal(t, v1.EndpointServiceID, v3.EndpointServiceID) }) } func Test_pickDefaultAZ(t *testing.T) { t.Run("normal", func(t *testing.T) { _, _, err := pickDefaultAZ() assert.NoError(t, err) }) } func TestAppendAssumeEntity(t *testing.T) { req := mockReq() iam, e := BuildAgentRoleRequirements(req) assert.NoError(t, e) // //t.Run("failed", func(t *testing.T) { // _, _, err := AppendAssumeEntity(context.TODO(), iam, "useraaaa") // assert.Error(t, err) //}) t.Run("not updated", func(t *testing.T) { copyIAM := *iam newIAM, updated, err := AppendAssumeEntity(context.TODO(), ©IAM, "useraaaa") assert.NoError(t, err) assert.False(t, updated) assert.Equal(t, iam.Spec.Aws.AssumeRolePolicyDocument, newIAM.Spec.Aws.AssumeRolePolicyDocument) }) t.Run("updated", func(t *testing.T) { copyIAM := *iam newIAM, updated, err := AppendAssumeEntity(context.TODO(), ©IAM, "newarn") assert.NoError(t, err) assert.True(t, updated) assert.NotEqual(t, iam.Spec.Aws.AssumeRolePolicyDocument, newIAM.Spec.Aws.AssumeRolePolicyDocument) }) } func TestAppendAssumeEntityGCP(t *testing.T) { req := mockReq() origin := req.Vendor originArn := req.AgentAccessARN req.AgentAccessARN = "useraaaa@gcp-autonomous-dev.iam.gserviceaccount.com" defer func() { req.Vendor = origin req.AgentAccessARN = originArn }() originalService := configs.GetServiceInfo() defer func() { configs.Init(originalService) }() configs.Init(structs.ServiceInfo{ ServerName: constants.ObservabilityRegionalServer, Port: 8080, Vendor: string(constants.GCP), Region: defaultRegion, Host: "127.0.0.1", Env: constants.EnvDev, MonitorSpec: map[string]structs.MonitorSpec{ "default": { Cpu: "10m", Memory: "200Mi", InitialStorage: "40Gi", OverSell: 1, IsDefault: true, }, }, AZConfig: map[string]structs.AvailabilityZoneConfig{ "zone1": { EndpointServiceID: "e1", IngressClassName: "i1", }, "zone2": { EndpointServiceID: "e2", IngressClassName: "i2", }, }, AccountInfo: structs.AccountInfoConfig{ Oidc: "oidc.eks.us-west-2.amazonaws.com/id/005D92AA81DE2DCE14A2991BAC54FEF8", OidcARN: "arn:aws:iam::552185537300:oidc-provider/oidc.eks.us-west-2.amazonaws.com/id/005D92AA81DE2DCE14A2991BAC54FEF8", Account: "aaaa", Region: "us-west-2", }, BizTypeConfigs: map[string]structs.BizTypeConfig{ "b": { AgentManagerARN: "aaa", }, "tidbcloud": { AgentManagerARN: "role2", }, }, }) iam, e := BuildAgentRoleRequirements(req) assert.NoError(t, e) t.Run("not updated", func(t *testing.T) { copyIAM := *iam newIAM, updated, err := AppendAssumeEntity(context.TODO(), ©IAM, "useraaaa@gcp-autonomous-dev.iam.gserviceaccount.com") assert.NoError(t, err) assert.False(t, updated) assert.Equal(t, iam.Spec.Gcp.AssumeRoleAccounts, newIAM.Spec.Gcp.AssumeRoleAccounts) }) t.Run("updated", func(t *testing.T) { copyIAM := *iam newIAM, updated, err := AppendAssumeEntity(context.TODO(), ©IAM, "newarn@gcp-autonomous-dev.iam.gserviceaccount.com") assert.NoError(t, err) assert.True(t, updated) assert.Equal(t, len(iam.Spec.Gcp.AssumeRoleAccounts)+1, len(newIAM.Spec.Gcp.AssumeRoleAccounts)) assert.NotEqual(t, len(iam.Spec.Gcp.AssumeRoleAccounts), len(newIAM.Spec.Gcp.AssumeRoleAccounts)) }) t.Run("wrong arn", func(t *testing.T) { copyIAM := *iam newIAM, updated, err := AppendAssumeEntity(context.TODO(), ©IAM, "newarn") assert.Error(t, err) assert.False(t, updated) assert.Nil(t, newIAM) }) } func TestAppendAssumeEntityAliCloud(t *testing.T) { req := mockReq() origin := req.Vendor originArn := req.AgentAccessARN req.AgentAccessARN = "useraaaa" defer func() { req.Vendor = origin req.AgentAccessARN = originArn }() originalService := configs.GetServiceInfo() defer func() { configs.Init(originalService) }() configs.Init(structs.ServiceInfo{ ServerName: constants.ObservabilityRegionalServer, Port: 8080, Vendor: string(constants.AliCloud), Region: defaultRegion, Host: "127.0.0.1", Env: constants.EnvDev, MonitorSpec: map[string]structs.MonitorSpec{ "default": { Cpu: "10m", Memory: "200Mi", InitialStorage: "40Gi", OverSell: 1, IsDefault: true, }, }, AZConfig: map[string]structs.AvailabilityZoneConfig{ "zone1": { EndpointServiceID: "e1", IngressClassName: "i1", }, "zone2": { EndpointServiceID: "e2", IngressClassName: "i2", }, }, AccountInfo: structs.AccountInfoConfig{ Oidc: "oidc.eks.us-west-2.amazonaws.com/id/005D92AA81DE2DCE14A2991BAC54FEF8", OidcARN: "arn:aws:iam::552185537300:oidc-provider/oidc.eks.us-west-2.amazonaws.com/id/005D92AA81DE2DCE14A2991BAC54FEF8", Account: "aaaa", Region: "us-west-2", }, BizTypeConfigs: map[string]structs.BizTypeConfig{ "b": { AgentManagerARN: "aaa", }, "tidbcloud": { AgentManagerARN: "role2", }, }, }) iam, e := BuildAgentRoleRequirements(req) assert.NoError(t, e) t.Run("not updated", func(t *testing.T) { copyIAM := *iam newIAM, updated, err := AppendAssumeEntity(context.TODO(), ©IAM, "useraaaa") assert.NoError(t, err) assert.False(t, updated) assert.Equal(t, iam.Spec.AliCloud.AssumeRolePolicyDocument, newIAM.Spec.AliCloud.AssumeRolePolicyDocument) }) t.Run("updated", func(t *testing.T) { copyIAM := *iam newIAM, updated, err := AppendAssumeEntity(context.TODO(), ©IAM, "newarn") assert.NoError(t, err) assert.True(t, updated) assert.NotEqual(t, iam.Spec.AliCloud.AssumeRolePolicyDocument, newIAM.Spec.AliCloud.AssumeRolePolicyDocument) }) } func TestBuildPrometheusRuleResourceRequirements(t *testing.T) { t.Run("normal", func(t *testing.T) { req := mockReq() req.MonitorPromRules.Groups = []prometheusv1.RuleGroup{ 0: { Name: "example", Interval: "5", Rules: []prometheusv1.Rule{ { Alert: "HighRequestLatency", Expr: intstr.FromString("request_latency_seconds:mean5m{job=\"myjob\"} > 0.5"), For: "10m", Labels: map[string]string{ "severity": "critical", }, Annotations: map[string]string{ "summary": "High request latency", }, }, }, }, } promRule, err := BuildPrometheusRuleResourceRequirements(req, &v1.Cluster{}) assert.NoError(t, err) assert.Equal(t, req.ID, promRule.Name) assert.Empty(t, promRule.Annotations) assert.NotEmpty(t, promRule.Labels) assert.NotEmpty(t, promRule.Spec.Groups) assert.Equal(t, "tenant111", promRule.Labels[constants.LabelTenantID]) }) t.Run("build vm rule", func(t *testing.T) { original := configs.GetServiceInfo() defer configs.Init(original) configs.Init(structs.ServiceInfo{ BizTypeConfigs: map[string]structs.BizTypeConfig{ "b": { BizType: "b", MultiTenancyRule: true, }, }, }) req := mockReq() req.ShareMode = constants.SHARED req.BizType = "b" req.MuteAlert = true req.MonitorPromRules.Groups = []prometheusv1.RuleGroup{ 0: { Name: "example", Interval: "5", Rules: []prometheusv1.Rule{ { Alert: "HighRequestLatency", Expr: intstr.FromString("request_latency_seconds:mean5m{job=\"myjob\"} > 0.5"), For: "10m", Labels: map[string]string{ "severity": "critical", }, Annotations: map[string]string{ "summary": "High request latency", }, }, }, }, } promRule, err := BuildPrometheusRuleResourceRequirements(req, &v1.Cluster{ Spec: v1.ClusterSpec{ VictoriaAccount: v1.VictoriaAccount{ AccountID: 1001, ProjectID: 0, }, }, }) assert.NoError(t, err) assert.Equal(t, req.ID, promRule.Name) assert.Empty(t, promRule.Annotations) assert.NotEmpty(t, promRule.Labels) assert.NotEmpty(t, promRule.Spec.Groups) assert.Equal(t, "tenant111", promRule.Labels[constants.LabelTenantID]) assert.Empty(t, promRule.ObjectMeta.Labels[constants.LabelVMAccountID]) assert.Empty(t, promRule.ObjectMeta.Labels[constants.LabelVMProjectID]) assert.Empty(t, promRule.ObjectMeta.Labels[constants.LabelMultiTenancyRule]) }) t.Run("build vm rule and muteAlert=false", func(t *testing.T) { original := configs.GetServiceInfo() defer configs.Init(original) configs.Init(structs.ServiceInfo{ BizTypeConfigs: map[string]structs.BizTypeConfig{ "b": { BizType: "b", MultiTenancyRule: true, }, }, }) req := mockReq() req.ShareMode = constants.SHARED req.BizType = "b" req.MonitorPromRules.Groups = []prometheusv1.RuleGroup{ 0: { Name: "example", Interval: "5", Rules: []prometheusv1.Rule{ { Alert: "HighRequestLatency", Expr: intstr.FromString("request_latency_seconds:mean5m{job=\"myjob\"} > 0.5"), For: "10m", Labels: map[string]string{ "severity": "critical", }, Annotations: map[string]string{ "summary": "High request latency", }, }, }, }, } promRule, err := BuildPrometheusRuleResourceRequirements(req, &v1.Cluster{ Spec: v1.ClusterSpec{ VictoriaAccount: v1.VictoriaAccount{ AccountID: 1001, ProjectID: 0, }, }, }) assert.NoError(t, err) assert.Equal(t, req.ID, promRule.Name) assert.Empty(t, promRule.Annotations) assert.NotEmpty(t, promRule.Labels) assert.NotEmpty(t, promRule.Spec.Groups) assert.Equal(t, "tenant111", promRule.Labels[constants.LabelTenantID]) assert.Equal(t, "1001", promRule.ObjectMeta.Labels[constants.LabelVMAccountID]) assert.Equal(t, "0", promRule.ObjectMeta.Labels[constants.LabelVMProjectID]) assert.Equal(t, "true", promRule.ObjectMeta.Labels[constants.LabelMultiTenancyRule]) }) } func TestGetAlertingConfig(t *testing.T) { original := configs.GetServiceInfo() defer configs.Init(original) configs.Init(structs.ServiceInfo{ AlertManagerInfo: structs.AlertManagerConfig{ Name: "alertmanager-operated", NameSpace: "o11y", Port: "web", }, }) alertConfig := getAlertingConfig() assert.NotEmpty(t, alertConfig.Alertmanagers) assert.Equal(t, 1, len(alertConfig.Alertmanagers)) assert.Equal(t, "alertmanager-operated", alertConfig.Alertmanagers[0].Name) assert.Equal(t, "o11y", alertConfig.Alertmanagers[0].Namespace) assert.Equal(t, intstr.FromString("web"), alertConfig.Alertmanagers[0].Port) } func TestBuildPrometheusRuleTemplateResourceRequirements(t *testing.T) { t.Run("normal", func(t *testing.T) { req := message.ModifyAlertRulesTemplateReq{} req.TemplateID = "test" req.MonitorPromRules.Groups = []prometheusv1.RuleGroup{ 0: { Name: "example", Interval: "5", Rules: []prometheusv1.Rule{ { Alert: "HighRequestLatency", Expr: intstr.FromString("request_latency_seconds:mean5m{job=\"myjob\"} > 0.5"), For: "10m", Labels: map[string]string{ "severity": "critical", }, Annotations: map[string]string{ "summary": "High request latency", }, }, }, }, } promRule, err := BuildPrometheusRuleTemplateResourceRequirements(req) assert.NoError(t, err) assert.Equal(t, "test", promRule.Name) assert.Equal(t, constants.AlertRuleTemplateNamespace, promRule.Namespace) }) } func Test_getVMSpace(t *testing.T) { // Helper function to create base request createBaseReq := func(parentClusterID, keySpaceID, applicationID string) message.CreateClusterReq { return message.CreateClusterReq{ ClusterBaseInfo: structs.ClusterBaseInfo{ CommonReqInfo: structs.CommonReqInfo{ Vendor: constants.AWS, Region: "us-east-1", }, TenantID: "tenant1", ProjectID: "project1", ApplicationID: applicationID, BizType: "tidbcloud", ID: "cluster1", }, ParentClusterID: parentClusterID, KeySpaceID: keySpaceID, } } tests := []struct { name string req message.CreateClusterReq setupMocks func(*mockresourcemanager.MockClusterReaderWriter, *mockresourcemanager.MockDataAccountReaderWriter, *mockresourcemanager.MockArchivedClusterReaderWriter) expectedAccountID int32 expectedProjectID int32 expectError bool }{ { name: "with parent cluster", req: createBaseReq("parent123", "456", "app1"), setupMocks: func(clusterRW *mockresourcemanager.MockClusterReaderWriter, dataAccountRW *mockresourcemanager.MockDataAccountReaderWriter, archivedClusterRW *mockresourcemanager.MockArchivedClusterReaderWriter) { clusterRW.EXPECT().GetCluster(gomock.Any(), "parent123").Return(&v1.Cluster{ Spec: v1.ClusterSpec{ VictoriaAccount: v1.VictoriaAccount{ AccountID: 1278, }, }, }, nil) }, expectedAccountID: 1278, expectedProjectID: 456, expectError: false, }, { name: "with parent cluster empty keyspace", req: createBaseReq("parent123", "", "app1"), setupMocks: func(clusterRW *mockresourcemanager.MockClusterReaderWriter, dataAccountRW *mockresourcemanager.MockDataAccountReaderWriter, archivedClusterRW *mockresourcemanager.MockArchivedClusterReaderWriter) { clusterRW.EXPECT().GetCluster(gomock.Any(), "parent123").Return(&v1.Cluster{ Spec: v1.ClusterSpec{ VictoriaAccount: v1.VictoriaAccount{ AccountID: 1278, }, }, }, nil) }, expectedAccountID: 1278, expectedProjectID: 0, expectError: false, }, { name: "parent cluster not found", req: createBaseReq("parent123", "456", "app1"), setupMocks: func(clusterRW *mockresourcemanager.MockClusterReaderWriter, dataAccountRW *mockresourcemanager.MockDataAccountReaderWriter, archivedClusterRW *mockresourcemanager.MockArchivedClusterReaderWriter) { clusterRW.EXPECT().GetCluster(gomock.Any(), "parent123").Return(nil, fmt.Errorf("cluster not found")) }, expectedAccountID: 0, expectedProjectID: 0, expectError: true, }, { name: "no parent cluster - k8s-infra with existing vmaccount", req: createBaseReq("", "789", constants.ApplicationIDK8sInfra), setupMocks: func(clusterRW *mockresourcemanager.MockClusterReaderWriter, dataAccountRW *mockresourcemanager.MockDataAccountReaderWriter, archivedClusterRW *mockresourcemanager.MockArchivedClusterReaderWriter) { archivedClusterRW.EXPECT().Get(gomock.Any(), gomock.Any(), gomock.Any()).Return(structs.ArchivedClusterInfo{ VMAccountID: "5555", }, nil) }, expectedAccountID: 5555, expectedProjectID: 0, // Function returns early, projectID is not set expectError: false, }, { name: "no parent cluster - k8s-infra without existing vmaccount", req: createBaseReq("", "789", constants.ApplicationIDK8sInfra), setupMocks: func(clusterRW *mockresourcemanager.MockClusterReaderWriter, dataAccountRW *mockresourcemanager.MockDataAccountReaderWriter, archivedClusterRW *mockresourcemanager.MockArchivedClusterReaderWriter) { archivedClusterRW.EXPECT().Get(gomock.Any(), gomock.Any(), gomock.Any()).Return(structs.ArchivedClusterInfo{ VMAccountID: "", }, nil) dataAccountRW.EXPECT().UniqueInt32ForCluster(gomock.Any()).Return(int32(3333), nil) }, expectedAccountID: 3333, expectedProjectID: 789, expectError: false, }, { name: "no parent cluster - k8s-infra archived cluster error", req: createBaseReq("", "789", constants.ApplicationIDK8sInfra), setupMocks: func(clusterRW *mockresourcemanager.MockClusterReaderWriter, dataAccountRW *mockresourcemanager.MockDataAccountReaderWriter, archivedClusterRW *mockresourcemanager.MockArchivedClusterReaderWriter) { archivedClusterRW.EXPECT().Get(gomock.Any(), gomock.Any(), gomock.Any()).Return(structs.ArchivedClusterInfo{}, fmt.Errorf("archived cluster error")) }, expectedAccountID: 0, expectedProjectID: 0, expectError: true, }, { name: "no parent cluster - regular cluster", req: createBaseReq("", "123", "regular-app"), setupMocks: func(clusterRW *mockresourcemanager.MockClusterReaderWriter, dataAccountRW *mockresourcemanager.MockDataAccountReaderWriter, archivedClusterRW *mockresourcemanager.MockArchivedClusterReaderWriter) { dataAccountRW.EXPECT().UniqueInt32ForCluster(gomock.Any()).Return(int32(4444), nil) }, expectedAccountID: 4444, expectedProjectID: 123, expectError: false, }, { name: "invalid keyspace id", req: createBaseReq("", "invalid", "regular-app"), setupMocks: func(clusterRW *mockresourcemanager.MockClusterReaderWriter, dataAccountRW *mockresourcemanager.MockDataAccountReaderWriter, archivedClusterRW *mockresourcemanager.MockArchivedClusterReaderWriter) { dataAccountRW.EXPECT().UniqueInt32ForCluster(gomock.Any()).Return(int32(4444), nil) }, expectedAccountID: 4444, // AccountID is generated before keyspace parsing fails expectedProjectID: 0, expectError: true, }, { name: "vmaccount id parse error", req: createBaseReq("", "123", constants.ApplicationIDK8sInfra), setupMocks: func(clusterRW *mockresourcemanager.MockClusterReaderWriter, dataAccountRW *mockresourcemanager.MockDataAccountReaderWriter, archivedClusterRW *mockresourcemanager.MockArchivedClusterReaderWriter) { archivedClusterRW.EXPECT().Get(gomock.Any(), gomock.Any(), gomock.Any()).Return(structs.ArchivedClusterInfo{ VMAccountID: "invalid-number", }, nil) }, expectedAccountID: 0, expectedProjectID: 0, expectError: true, }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { ctrl := gomock.NewController(t) defer ctrl.Finish() clusterRW := mockresourcemanager.NewMockClusterReaderWriter(ctrl) resourcemanager.MockClusterReaderWriter(clusterRW) dataAccountRW := mockresourcemanager.NewMockDataAccountReaderWriter(ctrl) resourcemanager.MockDataAccountReaderWriter(dataAccountRW) archivedClusterRW := mockresourcemanager.NewMockArchivedClusterReaderWriter(ctrl) resourcemanager.MockArchivedClusterReaderWriter(archivedClusterRW) // Setup mocks for this test case if tt.setupMocks != nil { tt.setupMocks(clusterRW, dataAccountRW, archivedClusterRW) } // Execute the function accountID, projectID, err := getVMSpace(context.TODO(), tt.req) // Verify results if tt.expectError { assert.Error(t, err) } else { assert.NoError(t, err) } assert.Equal(t, tt.expectedAccountID, accountID) assert.Equal(t, tt.expectedProjectID, projectID) }) } } var test_shared_policy = `{ "Version": "2012-10-17", "Statement": [ { "Effect": "Allow", "Action": [ "s3:ListBucket" ], "Resource": "arn:aws:s3:::testbucket" }, { "Effect": "Allow", "Action": [ "s3:*" ], "Resource": "arn:aws:s3:::testbucket/*" } ] }` func Test_buildPolicyDocument(t *testing.T) { t.Run("shared", func(t *testing.T) { result := buildPolicyDocument("testbucket", message.CreateClusterReq{ AgentSharedRole: true, }) assert.Equal(t, test_shared_policy, result) }) } func Test_buildGcpPolicies(t *testing.T) { t.Run("always shared regardless of agentSharedRole", func(t *testing.T) { policyA := buildGcpPolicies("testbucket", message.CreateClusterReq{ AgentSharedRole: true, ClusterBaseInfo: structs.ClusterBaseInfo{ TenantID: "tenant-a", ID: "cluster-a", }, }) policyB := buildGcpPolicies("testbucket", message.CreateClusterReq{ AgentSharedRole: false, ClusterBaseInfo: structs.ClusterBaseInfo{ TenantID: "tenant-b", ID: "cluster-b", }, }) key := v1.BuildGCPPolicyKey(v1.GCPResourceTypeCloudStorage, "testbucket") assert.Equal(t, 1, len(policyA[key])) assert.Equal(t, 1, len(policyB[key])) assert.Equal(t, "", policyA[key][0].PolicyCondition) assert.Equal(t, "", policyB[key][0].PolicyCondition) }) } func TestAppendTrustPolicyStatement(t *testing.T) { t.Run("oidc exists", func(t *testing.T) { role, err := generateRole() assert.NoError(t, err) req := message.CreateClusterReq{ AgentEnvironment: &structs.AgentEnvironment{ AWS: &structs.AWSAgentEnvironment{ Account: "testAcc", OIDC: "testOidc", }, }, } iam, update, err := AppendTrustPolicyStatement(context.Background(), role.DeepCopy(), req) assert.NoError(t, err) assert.Equal(t, false, update) assert.Equal(t, role.Spec, iam.Spec) }) t.Run("oidc not exists", func(t *testing.T) { role, err := generateRole() assert.NoError(t, err) req := message.CreateClusterReq{ AgentEnvironment: &structs.AgentEnvironment{ AWS: &structs.AWSAgentEnvironment{ Account: "testAcc", OIDC: "testOidc-1", }, }, } iam, update, err := AppendTrustPolicyStatement(context.Background(), role.DeepCopy(), req) assert.NoError(t, err) assert.Equal(t, update, true) assert.NotEqual(t, role.Spec.Aws.AssumeRolePolicyDocument, iam.Spec.Aws.AssumeRolePolicyDocument) }) } func generateRole() (*v1.IAM, error) { role := &v1.IAM{} policy := createDataPlaneAssumeRolePolicy("testAcc", "testOidc") policyJson, err := json.Marshal(policy) if err != nil { return nil, err } role.Spec.Aws.AssumeRolePolicyDocument = base64.StdEncoding.EncodeToString(policyJson) return role, nil } func generateAliCloudRole() (*v1.IAM, error) { role := &v1.IAM{} policy := createDataPlaneAssumeRolePolicyAliCloud("testAcc", "testOidc") policyJson, err := json.Marshal(policy) if err != nil { return nil, err } role.Spec.AliCloud.AssumeRolePolicyDocument = base64.StdEncoding.EncodeToString(policyJson) return role, nil } func TestAppendDataPlaneAssumeEntity(t *testing.T) { t.Run("AWS - new entity", func(t *testing.T) { role, err := generateRole() assert.NoError(t, err) newIAM, update, err := AppendDataPlaneAssumeEntity(context.Background(), role, mockSharedReq(withAgentAWS())) assert.NoError(t, err) assert.True(t, update) assert.NotEqual(t, role.Spec.Aws.AssumeRolePolicyDocument, newIAM.Spec.Aws.AssumeRolePolicyDocument) var policy map[string]interface{} decodedPolicy, _ := base64.StdEncoding.DecodeString(newIAM.Spec.Aws.AssumeRolePolicyDocument) json.Unmarshal(decodedPolicy, &policy) statements := policy["Statement"].([]interface{}) assert.Equal(t, 2, len(statements)) }) t.Run("GCP - new entity", func(t *testing.T) { originalService := configs.GetServiceInfo() defer func() { configs.Init(originalService) }() configs.Init(structs.ServiceInfo{ Vendor: string(constants.GCP), }) role := &v1.IAM{ Spec: v1.IAMSpec{ Gcp: v1.GCPIAM{ AssumeRoleAccounts: []string{"existing-account@example.com"}, }, }, } newIAM, update, err := AppendDataPlaneAssumeEntity(context.Background(), role, mockSharedReq(withAgentGCP())) assert.NoError(t, err) assert.True(t, update) assert.Equal(t, 2, len(newIAM.Spec.Gcp.AssumeRoleAccounts)) }) t.Run("Azure - new entity", func(t *testing.T) { role := &v1.IAM{ Spec: v1.IAMSpec{ Azure: &v1.AzureIAM{ WorkloadServiceAccounts: []string{"existing-account"}, }, }, } originalService := configs.GetServiceInfo() defer func() { configs.Init(originalService) }() configs.Init(structs.ServiceInfo{ Vendor: string(constants.Azure), }) newIAM, update, err := AppendDataPlaneAssumeEntity(context.Background(), role, mockSharedReq(withAgentAzure())) assert.NoError(t, err) assert.True(t, update) assert.Equal(t, 2, len(newIAM.Spec.Azure.WorkloadServiceAccounts)) }) t.Run("AliCloud - new entity", func(t *testing.T) { originalService := configs.GetServiceInfo() defer func() { configs.Init(originalService) }() configs.Init(structs.ServiceInfo{ Vendor: string(constants.AliCloud), }) role, err := generateAliCloudRole() assert.NoError(t, err) req := mockSharedReq(withAgentAliCloud()) req.Vendor = "alicloud" newIAM, update, err := AppendDataPlaneAssumeEntity(context.Background(), role, req) assert.NoError(t, err) assert.True(t, update) assert.NotEqual(t, role.Spec.AliCloud.AssumeRolePolicyDocument, newIAM.Spec.AliCloud.AssumeRolePolicyDocument) var policy map[string]interface{} decodedPolicy, _ := base64.StdEncoding.DecodeString(newIAM.Spec.AliCloud.AssumeRolePolicyDocument) json.Unmarshal(decodedPolicy, &policy) statements := policy["Statement"].([]interface{}) assert.Equal(t, 2, len(statements)) }) } func TestDataPlaneSANamespaceName(t *testing.T) { t.Run("GCP with k8s-infra application", func(t *testing.T) { result := DataPlaneSANamespaceName(constants.GCP, "k8s-infra") assert.Equal(t, "monitoring/o11y-vector", result) }) t.Run("GCP with custom application", func(t *testing.T) { result := DataPlaneSANamespaceName(constants.GCP, "application111") assert.Equal(t, "tidbapplication111/o11y-vector-agent", result) }) t.Run("Azure with k8s-infra application", func(t *testing.T) { result := DataPlaneSANamespaceName(constants.Azure, "k8s-infra") assert.Equal(t, "monitoring:o11y-vector", result) }) t.Run("Azure with custom application", func(t *testing.T) { result := DataPlaneSANamespaceName(constants.Azure, "application222") assert.Equal(t, "tidbapplication222:o11y-vector-agent", result) }) t.Run("AWS vendor returns empty string", func(t *testing.T) { result := DataPlaneSANamespaceName(constants.AWS, "application333") assert.Equal(t, "", result) }) t.Run("AliCloud vendor returns empty string", func(t *testing.T) { result := DataPlaneSANamespaceName(constants.AliCloud, "application444") assert.Equal(t, "", result) }) t.Run("Unknown vendor returns empty string", func(t *testing.T) { result := DataPlaneSANamespaceName("unknown", "application555") assert.Equal(t, "", result) }) t.Run("Empty application ID", func(t *testing.T) { result := DataPlaneSANamespaceName(constants.GCP, "") assert.Equal(t, "tidb/o11y-vector-agent", result) result = DataPlaneSANamespaceName(constants.Azure, "") assert.Equal(t, "tidb:o11y-vector-agent", result) }) } func Test_appendVM(t *testing.T) { ctrl := gomock.NewController(t) defer ctrl.Finish() dataAccountRW := mockresourcemanager.NewMockDataAccountReaderWriter(ctrl) resourcemanager.MockDataAccountReaderWriter(dataAccountRW) clusterRW := mockresourcemanager.NewMockClusterReaderWriter(ctrl) resourcemanager.MockClusterReaderWriter(clusterRW) t.Run("legacy mode - not nextgen", func(t *testing.T) { clusterRW.EXPECT().GetCluster(gomock.Any(), "parent123").Return(&v1.Cluster{ Spec: v1.ClusterSpec{ VictoriaAccount: v1.VictoriaAccount{ AccountID: 7777, }, }, }, nil).Times(1) req := mockSharedReq() req.AgentNextGen = false req.ParentClusterID = "parent123" req.KeySpaceID = "123" cluster := &v1.Cluster{ ObjectMeta: metav1.ObjectMeta{ Name: req.ID, Namespace: NamespaceName(req.BizType, req.TenantID), }, } err := appendVM(context.TODO(), cluster, req, "zone2", "ignored") assert.NoError(t, err) assert.Equal(t, int32(7777), cluster.Spec.VictoriaAccount.AccountID) // From parent assert.Equal(t, int32(123), cluster.Spec.VictoriaAccount.ProjectID) // From KeySpaceID assert.Equal(t, "zone2", cluster.Spec.VictoriaAccount.AvailabilityZone) }) }