package networkmetering import ( "context" "errors" "os" "testing" "github.com/stretchr/testify/assert" "github.com/tidbcloud/metering/pkg/common/config" "github.com/tidbcloud/metering/pkg/common/constants" "github.com/tidbcloud/metering/pkg/common/models" "github.com/tidbcloud/metering/test/mockapigateway" "github.com/tidbcloud/metering/test/mocknetworkmetering" "go.uber.org/mock/gomock" "go.uber.org/zap" ) func TestPickAzureTagsV2024(t *testing.T) { log, _ := zap.NewDevelopment() network := &NetworkMetering{ conf: &config.ServiceConfig{ Vendor: constants.VendorAzure, Env: constants.EnvDev, TrafficMeteringConfig: &config.TrafficMeteringConfig{ Provider: config.TrafficMeteringProvider{ Azure: &config.TrafficMeteringProviderAzureConfig{ BillingSchemaVersion: constants.AzureSchemaVersionV2024, }, }}, }, log: log, } source := `"aks-managed-consolidated-additional-properties": "b897a180-5f9c-11ef-ae65-9e4b0923f854","aks-managed-createOperationID": "fdd5f6d3-272e-4de4-be69-1356e2a0af0e","aks-managed-creationSource": "vmssclient-aks-n59d3c535a9d-70638116-vmss","aks-managed-kubeletIdentityClientID": "32259954-03b5-4cb8-89d1-e84fabea6c36","aks-managed-orchestrator": "Kubernetes:1.30.3","aks-managed-poolName": "n59d3c535a9d","aks-managed-resourceNameSuffix": "37478521","aks-managed-ssh-access": "LocalUser","cluster": "shared","usedby": "customer-tidb"` tags := network.PickAzureTagValue(source) assert.Equal(t, "b897a180-5f9c-11ef-ae65-9e4b0923f854", tags["aks-managed-consolidated-additional-properties"]) assert.Equal(t, "fdd5f6d3-272e-4de4-be69-1356e2a0af0e", tags["aks-managed-createOperationID"]) assert.Equal(t, "vmssclient-aks-n59d3c535a9d-70638116-vmss", tags["aks-managed-creationSource"]) assert.Equal(t, "32259954-03b5-4cb8-89d1-e84fabea6c36", tags["aks-managed-kubeletIdentityClientID"]) assert.Equal(t, "Kubernetes:1.30.3", tags["aks-managed-orchestrator"]) assert.Equal(t, "n59d3c535a9d", tags["aks-managed-poolName"]) assert.Equal(t, "37478521", tags["aks-managed-resourceNameSuffix"]) assert.Equal(t, "LocalUser", tags["aks-managed-ssh-access"]) assert.Equal(t, "shared", tags["cluster"]) assert.Equal(t, "customer-tidb", tags["usedby"]) } func TestPickAzureTagsV2021(t *testing.T) { log, _ := zap.NewDevelopment() network := &NetworkMetering{ conf: &config.ServiceConfig{ Vendor: constants.VendorAzure, Env: constants.EnvDev, TrafficMeteringConfig: &config.TrafficMeteringConfig{ Provider: config.TrafficMeteringProvider{ Azure: &config.TrafficMeteringProviderAzureConfig{ BillingSchemaVersion: constants.AzureSchemaVersionV2021, }, }}, }, log: log, } source := `{"aks-managed-bootstrap-profile-acr-name":"","aks-managed-consolidated-additional-properties":"852e83a3-67a6-11f0-adf2-c2e1fcecf750","aks-managed-coordination":"true","aks-managed-createOperationID":"2decd873-7072-432c-bd58-27b9a7b3f936","aks-managed-creationSource":"vmssclient-aks-n2295da69929-19159213-vmss","aks-managed-enable-imds-restriction":"false","aks-managed-kubeletIdentityClientID":"f5457795-1972-4a39-b773-f1cb51a1f8d0","aks-managed-orchestrator":"Kubernetes:1.32.5","aks-managed-poolName":"n2295da69929","aks-managed-resourceNameSuffix":"20660695","aks-managed-ssh-access":"LocalUser","cluster":"shared","project":"1369847559694040704","servicetype":"dedicated","tenant":"1369847559692509562","usedby":"customer-tidb"}` tags := network.PickAzureTagValue(source) assert.Equal(t, "852e83a3-67a6-11f0-adf2-c2e1fcecf750", tags["aks-managed-consolidated-additional-properties"]) assert.Equal(t, "2decd873-7072-432c-bd58-27b9a7b3f936", tags["aks-managed-createOperationID"]) assert.Equal(t, "vmssclient-aks-n2295da69929-19159213-vmss", tags["aks-managed-creationSource"]) assert.Equal(t, "f5457795-1972-4a39-b773-f1cb51a1f8d0", tags["aks-managed-kubeletIdentityClientID"]) assert.Equal(t, "Kubernetes:1.32.5", tags["aks-managed-orchestrator"]) assert.Equal(t, "n2295da69929", tags["aks-managed-poolName"]) assert.Equal(t, "20660695", tags["aks-managed-resourceNameSuffix"]) assert.Equal(t, "LocalUser", tags["aks-managed-ssh-access"]) assert.Equal(t, "shared", tags["cluster"]) assert.Equal(t, "customer-tidb", tags["usedby"]) } func TestBillingFloat2Int64High(t *testing.T) { assert.Equal(t, int64(0), billingFloat2Int64High(0)) assert.Equal(t, int64(10000000000), billingFloat2Int64High(1)) assert.Equal(t, int64(11000000000), billingFloat2Int64High(1.1)) assert.Equal(t, int64(15000000000), billingFloat2Int64High(1.5)) assert.Equal(t, int64(16000000000), billingFloat2Int64High(1.6)) } func TestBuildAWSUnitWithPremiumReservedSuffix(t *testing.T) { billing := &models.AWSBilling{ LineItemProductCode: "AmazonEC2", ProductServicename: "AWS Data Transfer", LineItemOperation: "InterZone-In", ProductRegion: "us-west-2", MeteringServicePlan: constants.ServicePlanPremiumReserved, } unit, ok := buildAWSUnit(billing) assert.True(t, ok) assert.Equal(t, "network-traffic:aws:us-west-2:reserved:cross_az", unit) } func TestBuildAWSUnitWithoutPremiumReservedSuffix(t *testing.T) { billing := &models.AWSBilling{ LineItemProductCode: "AmazonEC2", ProductServicename: "AWS Data Transfer", LineItemOperation: "InterZone-In", ProductRegion: "us-west-2", } unit, ok := buildAWSUnit(billing) assert.True(t, ok) assert.Equal(t, "network-traffic:aws:us-west-2:cross_az", unit) } func TestBuildAzureUnit2024(t *testing.T) { log, _ := zap.NewDevelopment() network := &NetworkMetering{ conf: &config.ServiceConfig{ Vendor: constants.VendorAzure, Env: constants.EnvDev, TrafficMeteringConfig: &config.TrafficMeteringConfig{ Provider: config.TrafficMeteringProvider{ Azure: &config.TrafficMeteringProviderAzureConfig{ BillingSchemaVersion: constants.AzureSchemaVersionV2024, }, }, }, }, log: log, } billing := &models.AzureBilling{ ConsumedService: "Microsoft.ContainerService/managedClusters", MeterCategory: "Bandwidth", ResourceLocation: "EastUS", } unit, skuType, ok := network.buildAzureUnit(billing) assert.True(t, ok) assert.Equal(t, "network-traffic:azure:eastus:internet", unit) assert.Equal(t, constants.TrafficInternet, skuType) billing = &models.AzureBilling{ ConsumedService: "Microsoft.ContainerService/managedClusters", MeterCategory: "Bandwidth", ResourceLocation: "EastUS", MeterSubCategory: "Inter-Region", } unit, skuType, ok = network.buildAzureUnit(billing) assert.True(t, ok) assert.Equal(t, "network-traffic:azure:eastus:cross_region", unit) assert.Equal(t, constants.TrafficCrossRegion, skuType) billing = &models.AzureBilling{ ConsumedService: "Microsoft.ContainerService/managedClusters", MeterCategory: "Load Balancer", ResourceLocation: "EastUS", } unit, skuType, ok = network.buildAzureUnit(billing) assert.True(t, ok) assert.Equal(t, "network-traffic:azure:eastus:lb", unit) assert.Equal(t, constants.TrafficLB, skuType) billing = &models.AzureBilling{ ConsumedService: "Microsoft.ContainerService/managedClusters", MeterCategory: "NAT Gateway", ResourceLocation: "EastUS", } unit, skuType, ok = network.buildAzureUnit(billing) assert.True(t, ok) assert.Equal(t, "network-traffic:azure:eastus:nat", unit) assert.Equal(t, constants.TrafficNAT, skuType) billing = &models.AzureBilling{ ConsumedService: "Microsoft.ContainerService/managedClusters", MeterCategory: "Virtual Network", ResourceLocation: "EastUS", MeterSubCategory: "IP Addresses", } unit, skuType, ok = network.buildAzureUnit(billing) assert.True(t, ok) assert.Equal(t, "network-traffic:azure:eastus:ipv4", unit) assert.Equal(t, constants.TrafficIpv4, skuType) } func TestBuildAzureUnit2021(t *testing.T) { log, _ := zap.NewDevelopment() network := &NetworkMetering{ conf: &config.ServiceConfig{ Vendor: constants.VendorAzure, Env: constants.EnvDev, TrafficMeteringConfig: &config.TrafficMeteringConfig{ Provider: config.TrafficMeteringProvider{ Azure: &config.TrafficMeteringProviderAzureConfig{ BillingSchemaVersion: constants.AzureSchemaVersionV2021, }, }, }, }, log: log, } billing := &models.AzureBilling{ ConsumedService: "Microsoft.ContainerService/managedClusters", MeterCategory: "Bandwidth", ResourceLocation: "EastUS", } unit, skuType, ok := network.buildAzureUnit(billing) assert.True(t, ok) assert.Equal(t, "network-traffic:azure:eastus:internet", unit) assert.Equal(t, constants.TrafficInternet, skuType) billing = &models.AzureBilling{ ConsumedService: "Microsoft.ContainerService/managedClusters", MeterCategory: "Bandwidth", ResourceLocation: "EastUS", MeterSubCategory: "Bandwidth Inter-Region", } unit, skuType, ok = network.buildAzureUnit(billing) assert.True(t, ok) assert.Equal(t, "network-traffic:azure:eastus:cross_region", unit) assert.Equal(t, constants.TrafficCrossRegion, skuType) // Test 2021 specific case: partial match for Inter-Region using contains billing = &models.AzureBilling{ ConsumedService: "Microsoft.ContainerService/managedClusters", MeterCategory: "Bandwidth", ResourceLocation: "EastUS", MeterSubCategory: "Data Transfer Inter-Region", } unit, skuType, ok = network.buildAzureUnit(billing) assert.True(t, ok) assert.Equal(t, "network-traffic:azure:eastus:cross_region", unit) assert.Equal(t, constants.TrafficCrossRegion, skuType) billing = &models.AzureBilling{ ConsumedService: "Microsoft.ContainerService/managedClusters", MeterCategory: "Load Balancer", ResourceLocation: "EastUS", } unit, skuType, ok = network.buildAzureUnit(billing) assert.True(t, ok) assert.Equal(t, "network-traffic:azure:eastus:lb", unit) assert.Equal(t, constants.TrafficLB, skuType) billing = &models.AzureBilling{ ConsumedService: "Microsoft.ContainerService/managedClusters", MeterCategory: "NAT Gateway", ResourceLocation: "EastUS", } unit, skuType, ok = network.buildAzureUnit(billing) assert.True(t, ok) assert.Equal(t, "network-traffic:azure:eastus:nat", unit) assert.Equal(t, constants.TrafficNAT, skuType) billing = &models.AzureBilling{ ConsumedService: "Microsoft.ContainerService/managedClusters", MeterCategory: "Virtual Network", ResourceLocation: "EastUS", MeterSubCategory: "IP Addresses", } unit, skuType, ok = network.buildAzureUnit(billing) assert.True(t, ok) assert.Equal(t, "network-traffic:azure:eastus:ipv4", unit) assert.Equal(t, constants.TrafficIpv4, skuType) } func TestDealBillingAzure(t *testing.T) { log, _ := zap.NewDevelopment() network := &NetworkMetering{ conf: &config.ServiceConfig{ Vendor: constants.VendorAzure, Env: constants.EnvDev, TrafficMeteringConfig: &config.TrafficMeteringConfig{ Provider: config.TrafficMeteringProvider{ Azure: &config.TrafficMeteringProviderAzureConfig{ BillingSchemaVersion: constants.AzureSchemaVersionV2024, }, }, }, }, log: log, } billings := []*models.AzureBilling{ { AccountOwnerId: "1", ConsumedService: "Microsoft.ContainerService/managedClusters", MeterCategory: "Bandwidth", ResourceLocation: "EastUS", ResourceId: "321", CostInBillingCurrency: 12.1, Tags: `"cluster": "shared","usedby": "customer-tidb","project": "123","tenant": "456","environment": "dev"`, }, { AccountOwnerId: "1", ConsumedService: "Microsoft.ContainerService/managedClusters", MeterCategory: "Bandwidth", ResourceLocation: "EastUS", ResourceId: "123", CostInBillingCurrency: 2.1, Tags: `"cluster": "123","usedby": "customer-tidb","project": "123","tenant": "456","environment": "dev"`, }, { AccountOwnerId: "1", ConsumedService: "Microsoft.ContainerService/managedClusters", MeterCategory: "Bandwidth", ResourceLocation: "EastUS", ResourceId: "123", CostInBillingCurrency: 3.9, Tags: `"cluster": "123","usedby": "customer-tidb","project": "123","tenant": "456","environment": "dev"`, }, { AccountOwnerId: "1", ConsumedService: "Microsoft.ContainerService/managedClusters", MeterCategory: "Virtual Network", MeterSubCategory: "IP Addresses", ResourceLocation: "EastUS", ResourceId: "1234", CostInBillingCurrency: 3.8, Tags: `"cluster": "123","project": "123","tenant": "456","environment": "dev","k8s-azure-service": "db-tidb-public"`, }, { AccountOwnerId: "1", ConsumedService: "Microsoft.ContainerService/managedClusters", MeterCategory: "Load Balancer", MeterSubCategory: "", ResourceLocation: "EastUS", ResourceId: "1234", CostInBillingCurrency: 5.8, Tags: `"project": "123","tenant": "456","environment": "dev","k8s-azure-service": "db-tidb-public"`, }, } metering := []*models.BillingMetering{ { ID: 1, TenantID: "456", ProjectID: "123", UsedBy: "customer-tidb", ResourceID: "1234", Cost: 2000000, BillingDate: 1723487507, Unit: "network-traffic:azure:eastus:lb", }, } // in this case, the result should be 1 //timestamp := 1723487507 , which utc is Mon Aug 12 2024 18:31:47 res, err := network.DealBillingAzure(billings, metering, 1723487507) assert.Nil(t, err) assert.Equal(t, 4, len(res)) for _, r := range res { if r.ID == 1 { assert.Equal(t, "network-traffic:azure:eastus:lb", r.Unit) assert.Equal(t, int64(580000000), r.Cost) break } } } func TestDealBillingRepoter(t *testing.T) { ctrl := gomock.NewController(t) gateway := mockapigateway.NewMockApiGateway(ctrl) gateway.EXPECT().CountBillableClusters(gomock.Any(), gomock.Any(), "123", gomock.Any(), gomock.Any()).Return(int64(2), nil).AnyTimes() gateway.EXPECT().CountBillableClusters(gomock.Any(), gomock.Any(), "321", gomock.Any(), gomock.Any()).Return(int64(0), nil).AnyTimes() log, _ := zap.NewDevelopment() network := &NetworkMetering{ gateway: gateway, conf: &config.ServiceConfig{ Vendor: constants.VendorAzure, }, log: log, } metering := []*models.BillingMetering{ { ID: 1, TenantID: "456", ProjectID: "123", ClusterID: "123", UsedBy: "customer-tidb", ResourceID: "123", Cost: 2000000, BillingDate: 1723487507, Unit: "network-traffic:azure:eastus:internet", }, { ID: 1, TenantID: "456", ProjectID: "123", ClusterID: "123", UsedBy: "customer-tidb", ResourceID: "3211", Cost: 2000000, BillingDate: 1723487507, Unit: "network-traffic:azure:eastus:internet", }, { ID: 1, TenantID: "456", ProjectID: "123", ClusterID: "123", UsedBy: "customer-tidb", ResourceID: "3211", Cost: 2000000, BillingDate: 1723487507, Unit: "network-traffic:azure:eastus:internet", }, { ID: 1, TenantID: "456", ProjectID: "123", ClusterID: "123", UsedBy: "customer-tidb", ResourceID: "1234", Cost: 2000000, BillingDate: 1723487507, Unit: "network-traffic:azure:eastus:cross_region", }, { ID: 1, TenantID: "456", ProjectID: "321", ClusterID: "123", UsedBy: "customer-tidb", ResourceID: "123", Cost: 2000000, BillingDate: 1723487507, Unit: "network-traffic:azure:eastus:internet", }, { ID: 1, TenantID: "456", ProjectID: "321", ClusterID: "123", UsedBy: "customer-tidb", ResourceID: "123", Cost: 2000000, BillingDate: 1723487507, Unit: "network-traffic:azure:eastus:internet", }, { ID: 1, TenantID: "456", ProjectID: "notnum", ClusterID: "123", UsedBy: "customer-tidb", ResourceID: "123", Cost: 2000000, BillingDate: 1723487507, Unit: "network-traffic:azure:eastus:internet", }, } repoters := []*models.BillingReport{ { TenantID: 456, CostID: "123", Cost: 4000000, CostTimestamp: 1723487507, Labels: `{"projectID":"123","clusterID":"123"}`, Unit: "network-traffic:azure:eastus:internet", CreatedAt: 1723487507, }, } // in this case, the result should be 1 //timestamp := 1723487507 , which utc is Mon Aug 12 2024 18:31:47 ctx := context.Background() ctx = context.WithValue(ctx, constants.CtxKeyRerunTsDiff, int64(2)) res, err := network.DealBillingRepoter(ctx, metering, repoters, 1723487507) assert.Nil(t, err) assert.Equal(t, 3, len(res)) } func TestDealBillingAws(t *testing.T) { log, _ := zap.NewDevelopment() network := &NetworkMetering{ conf: &config.ServiceConfig{ Vendor: constants.VendorAWS, Env: constants.EnvDev, }, log: log, } billings := []*models.AWSBilling{ { LineItemOperation: "LoadBalancing-PublicIP-", ProductRegion: "us-west-2", LineItemProductCode: "AWSELB", ResourceTagsTenant: "456", ResourceTagsProject: "123", ResourceTagsCluster: "123", ResourceTagsUsedby: "customer-tidb", LineItemResourceId: "123", LineItemUnblendedCost: 3.1, }, { LineItemOperation: "InterZone-", ProductRegion: "us-west-2", LineItemProductCode: "AmazonEC2", ProductServicename: "AWS Data Transfer", ResourceTagsTenant: "456", ResourceTagsProject: "123", ResourceTagsCluster: "123", ResourceTagsUsedby: "customer-tidb", LineItemResourceId: "123", LineItemUnblendedCost: 2.1, }, { LineItemOperation: "LoadBalancing:", ProductRegion: "us-west-2", LineItemProductCode: "AWSELB", PricingUnit: "Hrs", ResourceTagsComponent: "nlb", LineItemUsageAmount: 1, ResourceTagsTenant: "456", ResourceTagsProject: "123", ResourceTagsCluster: "123", ResourceTagsUsedby: "customer-tidb", LineItemResourceId: "123", LineItemUnblendedCost: 4.1, }, { LineItemOperation: "InterZone-", ProductRegion: "us-west-2", LineItemProductCode: "AmazonEC2", ProductServicename: "AWS Data Transfer", ResourceTagsTenant: "45678", ResourceTagsProject: "12345", ResourceTagsCluster: "jinweitest", ResourceTagsUsedby: "customer-tidb", LineItemResourceId: "123567", LineItemUnblendedCost: 2.1, }, } metering := []*models.BillingMetering{ { ID: 1, TenantID: "456", ProjectID: "123", ClusterID: "123", UsedBy: "customer-tidb", ResourceID: "123", Cost: 2000000, BillingDate: 1723487507, Unit: "network-traffic:aws:us-west-2:internet", }, } // in this case, the result should be 1 //timestamp := 1723487507 , which utc is Mon Aug 12 2024 18:31:47 res, err := network.DealBillingAWS(billings, metering, 1723487507) assert.Nil(t, err) assert.Equal(t, 5, len(res)) for _, r := range res { if r.TenantID == "45678" && r.ProjectID == "12345" { assert.Empty(t, r.ClusterID) } } } func TestDealBillingGcp(t *testing.T) { log, _ := zap.NewDevelopment() network := &NetworkMetering{ conf: &config.ServiceConfig{ Vendor: constants.VendorGCP, Env: constants.EnvDev, }, log: log, } billings := []*models.GCPBilling{ { ServiceName: "Compute Engine", Description: "Network Inter Zone", Region: "us-west1", UserLabelTenantID: "456", UserLabelProjectID: "123", UserLabelClusterID: "shared", UserLabelUsedBy: "customer-tidb", ResourceName: "123", Cost: 3.1, CostAtList: 3.5, }, { ServiceName: "Compute Engine", Description: "Network Internet", Region: "us-west1", UserLabelTenantID: "456", UserLabelProjectID: "123", UserLabelClusterID: "shared", UserLabelUsedBy: "customer-tidb", ResourceName: "123", Cost: 2.1, CostAtList: 2.5, }, { ServiceName: "Networking", Description: "Network Load Balancer", Region: "us-west1", UserLabelTenantID: "456", UserLabelProjectID: "123", UserLabelClusterID: "shared", UserLabelUsedBy: "customer-tidb", ResourceName: "123", Cost: 4.1, CostAtList: 4.5, }, } metering := []*models.BillingMetering{ { ID: 1, TenantID: "456", ProjectID: "123", ClusterID: "123", UsedBy: "customer-tidb", ResourceID: "123", Cost: 2000000, BillingDate: 1723487507, Unit: "network-traffic:gcp:us-west1:internet", }, } // in this case, the result should be 3 //timestamp := 1723487507 , which utc is Mon Aug 12 2024 18:31:47 res, err := network.DealBillingGCP(billings, metering, 1723487507) assert.Nil(t, err) assert.Equal(t, 3, len(res)) } func TestDealBillingGcpUseCostAtList(t *testing.T) { log, _ := zap.NewDevelopment() network := &NetworkMetering{ conf: &config.ServiceConfig{ Vendor: constants.VendorGCP, Env: constants.EnvDev, }, log: log, } billings := []*models.GCPBilling{ { ServiceName: "Compute Engine", Description: "Network Inter Zone", Region: "us-west1", UserLabelTenantID: "456", UserLabelProjectID: "123", UserLabelClusterID: "789", UserLabelUsedBy: "customer-tidb", ResourceName: "res1", Cost: 3.0, CostAtList: 5.0, }, { ServiceName: "Compute Engine", Description: "Network Inter Zone", Region: "us-west1", UserLabelTenantID: "456", UserLabelProjectID: "123", UserLabelClusterID: "789", UserLabelUsedBy: "customer-tidb", ResourceName: "res1", Cost: 2.0, CostAtList: 4.0, }, } // No existing metering, so cost should be purely from CostAtList res, err := network.DealBillingGCP(billings, nil, 1723487507) assert.Nil(t, err) assert.Equal(t, 1, len(res)) // CostAtList: 5.0 + 4.0 = 9.0, in 10^10 = 90000000000, /100 = 900000000 (10^8) assert.Equal(t, int64(900000000), res[0].Cost) } func TestFilterEnvironmentAzure(t *testing.T) { log, _ := zap.NewDevelopment() network := &NetworkMetering{ conf: &config.ServiceConfig{ Vendor: constants.VendorGCP, Env: constants.EnvDev, }, log: log, } assert.True(t, network.filterEnvironmentAzure("dedicated-dev-pd1", "")) } func TestNewNetworkMetering(t *testing.T) { os.Setenv("AWS_ACCESS_KEY_ID", "test") os.Setenv("AWS_SECRET_ACCESS_KEY", "test") defer os.Unsetenv("AWS_ACCESS_KEY_ID") defer os.Unsetenv("AWS_SECRET_ACCESS_KEY") conf := &config.ServiceConfig{ Vendor: "aws", Env: "dev", Secret: &config.Secret{ ApiGatewayKey: "test", }, TrafficMeteringConfig: &config.TrafficMeteringConfig{ RerunDays: 1, CachePath: "cachetest_init", BillingReportSwitchTime: 0, Provider: config.TrafficMeteringProvider{ AWS: &config.TrafficMeteringProviderAWSConfig{ StoreProvider: config.StoreProviderAWSConfig{ RoleArn: "arn:aws", Account: "123456789012", Region: "us-east-1", BucketName: "test-bucket", Perfix: "test-prefix", }, BillingProvider: []config.BillingProviderAWSConfig{{ RoleArn: "arn:aws", Account: "123456789012", Region: "us-east-1", BucketName: "test-bucket", Perfix: "test-prefix", SourceType: "csv", }, }, }, }, }, } log, _ := zap.NewDevelopment() _, err := NewNetworkMetering(conf, log) assert.Nil(t, err) } func TestMeteringAzure(t *testing.T) { ctrl := gomock.NewController(t) gateway := mockapigateway.NewMockApiGateway(ctrl) loader := mocknetworkmetering.NewMockNetworkBillingLoader(ctrl) store := mocknetworkmetering.NewMockNetworkBillingStorer(ctrl) conf := &config.ServiceConfig{ Vendor: "azure", Env: "dev", Secret: &config.Secret{ ApiGatewayKey: "test", }, TrafficMeteringConfig: &config.TrafficMeteringConfig{ RerunDays: 1, CachePath: "cachetest_init", BillingReportSwitchTime: 0, Provider: config.TrafficMeteringProvider{ Azure: &config.TrafficMeteringProviderAzureConfig{ BillingSchemaVersion: constants.AzureSchemaVersionV2024, }, }, }, } billings := []*models.AzureBilling{ { AccountOwnerId: "1", ConsumedService: "Microsoft.ContainerService/managedClusters", MeterCategory: "Bandwidth", ResourceLocation: "EastUS", ResourceId: "321", CostInBillingCurrency: 12.1, Tags: `"cluster": "shared","usedby": "customer-tidb","project": "123","tenant": "456","environment": "dev"`, }, { AccountOwnerId: "1", ConsumedService: "Microsoft.ContainerService/managedClusters", MeterCategory: "Bandwidth", ResourceLocation: "EastUS", ResourceId: "123", CostInBillingCurrency: 2.1, Tags: `"cluster": "123","usedby": "customer-tidb","project": "123","tenant": "456","environment": "dev"`, }, { AccountOwnerId: "1", ConsumedService: "Microsoft.ContainerService/managedClusters", MeterCategory: "Bandwidth", ResourceLocation: "EastUS", ResourceId: "123", CostInBillingCurrency: 3.9, Tags: `"cluster": "123","usedby": "customer-tidb","project": "123","tenant": "456","environment": "dev"`, }, { AccountOwnerId: "1", ConsumedService: "Microsoft.ContainerService/managedClusters", MeterCategory: "Virtual Network", MeterSubCategory: "IP Addresses", ResourceLocation: "EastUS", ResourceId: "1234", CostInBillingCurrency: 3.8, Tags: `"cluster": "123","usedby": "customer-tidb","project": "123","tenant": "456","environment": "dev","k8s-azure-service": "kubernetes-internal"`, }, } log, _ := zap.NewDevelopment() metering := &NetworkMetering{ gateway: gateway, loader: loader, store: store, conf: conf, log: log, } // for close //test LoadAzure failed loader.EXPECT().LoadAzure(gomock.Any(), gomock.Any()).Return(nil, errors.New("test")) _, _, err := metering.Metering(context.TODO(), 1723487507) assert.NotNil(t, err) //test LoadAzure errNometaFound loader.EXPECT().LoadAzure(gomock.Any(), gomock.Any()).Return(nil, errNometaFound) _, ok, err := metering.Metering(context.TODO(), 1723487507) assert.Nil(t, err) assert.False(t, ok) //test LoadAzure empty loader.EXPECT().LoadAzure(gomock.Any(), gomock.Any()).Return([]*models.AzureBilling{}, nil) _, ok, err = metering.Metering(context.TODO(), 1723487507) assert.Nil(t, err) assert.False(t, ok) // test load store failed store.EXPECT().CloseDB(gomock.Any(), gomock.Any()).AnyTimes() loader.EXPECT().LoadAzure(gomock.Any(), gomock.Any()).Return(billings, nil) store.EXPECT().Load(gomock.Any(), gomock.Any()).Return(nil, nil, errors.New("test")) _, ok, err = metering.Metering(context.TODO(), 1723487507) assert.NotNil(t, err) assert.False(t, ok) gateway.EXPECT().CountBillableClusters(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(int64(1), nil).AnyTimes() //store failed loader.EXPECT().LoadAzure(gomock.Any(), gomock.Any()).Return(billings, nil) store.EXPECT().Load(gomock.Any(), gomock.Any()).Return([]*models.BillingMetering{}, []*models.BillingReport{}, nil) store.EXPECT().Store(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(errors.New("test")) _, ok, err = metering.Metering(context.TODO(), 1723487507) assert.NotNil(t, err) assert.False(t, ok) //upload failed loader.EXPECT().LoadAzure(gomock.Any(), gomock.Any()).Return(billings, nil) store.EXPECT().Load(gomock.Any(), gomock.Any()).Return([]*models.BillingMetering{}, []*models.BillingReport{}, nil) store.EXPECT().Store(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil) store.EXPECT().Upload(gomock.Any(), gomock.Any()).Return("", errors.New("test")) _, ok, err = metering.Metering(context.TODO(), 1723487507) assert.NotNil(t, err) assert.False(t, ok) // metering success loader.EXPECT().LoadAzure(gomock.Any(), gomock.Any()).Return(billings, nil) store.EXPECT().Load(gomock.Any(), gomock.Any()).Return([]*models.BillingMetering{}, []*models.BillingReport{}, nil) store.EXPECT().Store(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil) store.EXPECT().Upload(gomock.Any(), gomock.Any()).Return("path", nil) path, ok, err := metering.Metering(context.TODO(), 1723487507) assert.Nil(t, err) assert.True(t, ok) assert.Equal(t, "path", path) //test not support vendor metering.conf.Vendor = "notsupport" _, _, err = metering.Metering(context.TODO(), 1723487507) assert.NotNil(t, err) } func TestMeteringAWS(t *testing.T) { ctrl := gomock.NewController(t) gateway := mockapigateway.NewMockApiGateway(ctrl) loader := mocknetworkmetering.NewMockNetworkBillingLoader(ctrl) store := mocknetworkmetering.NewMockNetworkBillingStorer(ctrl) conf := &config.ServiceConfig{ Vendor: "aws", Env: "dev", Secret: &config.Secret{ ApiGatewayKey: "test", }, TrafficMeteringConfig: &config.TrafficMeteringConfig{ RerunDays: 1, CachePath: "cachetest_init", BillingReportSwitchTime: 0, Provider: config.TrafficMeteringProvider{ AWS: &config.TrafficMeteringProviderAWSConfig{ StoreProvider: config.StoreProviderAWSConfig{ RoleArn: "arn:aws", Account: "123456789012", Region: "us-east-1", BucketName: "test-bucket", Perfix: "test-prefix", }, BillingProvider: []config.BillingProviderAWSConfig{{ RoleArn: "arn:aws", Account: "123456789012", Region: "us-east-1", BucketName: "test-bucket", Perfix: "test-prefix", SourceType: "csv", }, }, }, }, }, } billings := []*models.AWSBilling{ { LineItemOperation: "LoadBalancing-PublicIP-", ProductRegion: "us-west-2", LineItemProductCode: "AWSELB", ResourceTagsTenant: "456", ResourceTagsProject: "123", ResourceTagsCluster: "123", ResourceTagsUsedby: "customer-tidb", LineItemResourceId: "123", LineItemUnblendedCost: 3.1, }, { LineItemOperation: "InterZone-", ProductRegion: "us-west-2", LineItemProductCode: "AmazonEC2", ProductServicename: "AWS Data Transfer", ResourceTagsTenant: "456", ResourceTagsProject: "123", ResourceTagsCluster: "123", ResourceTagsUsedby: "customer-tidb", LineItemResourceId: "123", LineItemUnblendedCost: 2.1, }, { LineItemOperation: "LoadBalancing:", ProductRegion: "us-west-2", LineItemProductCode: "AWSELB", PricingUnit: "Hrs", ResourceTagsComponent: "nlb", LineItemUsageAmount: 1, ResourceTagsTenant: "456", ResourceTagsProject: "123", ResourceTagsCluster: "123", ResourceTagsUsedby: "customer-tidb", LineItemResourceId: "123", LineItemUnblendedCost: 4.1, }, { LineItemOperation: "InterZone-", ProductRegion: "us-west-2", LineItemProductCode: "AmazonEC2", ProductServicename: "AWS Data Transfer", ResourceTagsTenant: "45678", ResourceTagsProject: "12345", ResourceTagsCluster: "jinweitest", ResourceTagsUsedby: "customer-tidb", LineItemResourceId: "123567", LineItemUnblendedCost: 2.1, }, } log, _ := zap.NewDevelopment() metering := &NetworkMetering{ gateway: gateway, loader: loader, store: store, conf: conf, log: log, } //test LoadAWS failed loader.EXPECT().LoadAWS(gomock.Any(), gomock.Any()).Return(nil, errors.New("test")) _, _, err := metering.Metering(context.TODO(), 1723487507) assert.NotNil(t, err) //test LoadAWS errNometaFound loader.EXPECT().LoadAWS(gomock.Any(), gomock.Any()).Return(nil, errNometaFound) _, ok, err := metering.Metering(context.TODO(), 1723487507) assert.Nil(t, err) assert.False(t, ok) //test LoadAWS empty loader.EXPECT().LoadAWS(gomock.Any(), gomock.Any()).Return([]*models.AWSBilling{}, nil) _, ok, err = metering.Metering(context.TODO(), 1723487507) assert.Nil(t, err) assert.False(t, ok) // test load store failed store.EXPECT().CloseDB(gomock.Any(), gomock.Any()).AnyTimes() loader.EXPECT().LoadAWS(gomock.Any(), gomock.Any()).Return(billings, nil) store.EXPECT().Load(gomock.Any(), gomock.Any()).Return(nil, nil, errors.New("test")) _, ok, err = metering.Metering(context.TODO(), 1723487507) assert.NotNil(t, err) assert.False(t, ok) gateway.EXPECT().CountBillableClusters(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(int64(1), nil).AnyTimes() //store failed loader.EXPECT().LoadAWS(gomock.Any(), gomock.Any()).Return(billings, nil) store.EXPECT().Load(gomock.Any(), gomock.Any()).Return([]*models.BillingMetering{}, []*models.BillingReport{}, nil) store.EXPECT().Store(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(errors.New("test")) _, ok, err = metering.Metering(context.TODO(), 1723487507) assert.NotNil(t, err) assert.False(t, ok) //upload failed loader.EXPECT().LoadAWS(gomock.Any(), gomock.Any()).Return(billings, nil) store.EXPECT().Load(gomock.Any(), gomock.Any()).Return([]*models.BillingMetering{}, []*models.BillingReport{}, nil) store.EXPECT().Store(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil) store.EXPECT().Upload(gomock.Any(), gomock.Any()).Return("", errors.New("test")) _, ok, err = metering.Metering(context.TODO(), 1723487507) assert.NotNil(t, err) assert.False(t, ok) // metering success loader.EXPECT().LoadAWS(gomock.Any(), gomock.Any()).Return(billings, nil) store.EXPECT().Load(gomock.Any(), gomock.Any()).Return([]*models.BillingMetering{}, []*models.BillingReport{}, nil) store.EXPECT().Store(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil) store.EXPECT().Upload(gomock.Any(), gomock.Any()).Return("path", nil) path, ok, err := metering.Metering(context.TODO(), 1723487507) assert.Nil(t, err) assert.True(t, ok) assert.Equal(t, "path", path) } func TestMeteringGCP(t *testing.T) { ctrl := gomock.NewController(t) gateway := mockapigateway.NewMockApiGateway(ctrl) loader := mocknetworkmetering.NewMockNetworkBillingLoader(ctrl) store := mocknetworkmetering.NewMockNetworkBillingStorer(ctrl) conf := &config.ServiceConfig{ Vendor: "gcp", Env: "dev", Secret: &config.Secret{ ApiGatewayKey: "test", }, TrafficMeteringConfig: &config.TrafficMeteringConfig{ RerunDays: 1, CachePath: "cachetest_init", BillingReportSwitchTime: 0, Provider: config.TrafficMeteringProvider{ AWS: &config.TrafficMeteringProviderAWSConfig{ StoreProvider: config.StoreProviderAWSConfig{ RoleArn: "arn:aws", Account: "123456789012", Region: "us-east-1", BucketName: "test-bucket", Perfix: "test-prefix", }, BillingProvider: []config.BillingProviderAWSConfig{{ RoleArn: "arn:aws", Account: "123456789012", Region: "us-east-1", BucketName: "test-bucket", Perfix: "test-prefix", SourceType: "csv", }, }, }, }, }, } billings := []*models.GCPBilling{ { ServiceName: "Compute Engine", Description: "Network Inter Zone", Region: "us-west1", UserLabelTenantID: "456", UserLabelProjectID: "123", UserLabelClusterID: "shared", UserLabelUsedBy: "customer-tidb", ResourceName: "123", Cost: 3.1, }, { ServiceName: "Compute Engine", Description: "Network Internet", Region: "us-west1", UserLabelTenantID: "456", UserLabelProjectID: "123", UserLabelClusterID: "shared", UserLabelUsedBy: "customer-tidb", ResourceName: "123", Cost: 2.1, }, { ServiceName: "Networking", Description: "Network Load Balancer", Region: "us-west1", UserLabelTenantID: "456", UserLabelProjectID: "123", UserLabelClusterID: "shared", UserLabelUsedBy: "customer-tidb", ResourceName: "123", Cost: 4.1, }, } log, _ := zap.NewDevelopment() metering := &NetworkMetering{ gateway: gateway, loader: loader, store: store, conf: conf, log: log, } //test LoadGCP failed loader.EXPECT().LoadGCP(gomock.Any(), gomock.Any()).Return(nil, errors.New("test")) _, _, err := metering.Metering(context.TODO(), 1723487507) assert.NotNil(t, err) //test LoadGCP empty loader.EXPECT().LoadGCP(gomock.Any(), gomock.Any()).Return([]*models.GCPBilling{}, nil) _, ok, err := metering.Metering(context.TODO(), 1723487507) assert.Nil(t, err) assert.False(t, ok) // test load store failed store.EXPECT().CloseDB(gomock.Any(), gomock.Any()).AnyTimes() loader.EXPECT().LoadGCP(gomock.Any(), gomock.Any()).Return(billings, nil) store.EXPECT().Load(gomock.Any(), gomock.Any()).Return(nil, nil, errors.New("test")) _, ok, err = metering.Metering(context.TODO(), 1723487507) assert.NotNil(t, err) assert.False(t, ok) gateway.EXPECT().CountBillableClusters(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(int64(1), nil).AnyTimes() //store failed loader.EXPECT().LoadGCP(gomock.Any(), gomock.Any()).Return(billings, nil) store.EXPECT().Load(gomock.Any(), gomock.Any()).Return([]*models.BillingMetering{}, []*models.BillingReport{}, nil) store.EXPECT().Store(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(errors.New("test")) _, ok, err = metering.Metering(context.TODO(), 1723487507) assert.NotNil(t, err) assert.False(t, ok) //upload failed loader.EXPECT().LoadGCP(gomock.Any(), gomock.Any()).Return(billings, nil) store.EXPECT().Load(gomock.Any(), gomock.Any()).Return([]*models.BillingMetering{}, []*models.BillingReport{}, nil) store.EXPECT().Store(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil) store.EXPECT().Upload(gomock.Any(), gomock.Any()).Return("", errors.New("test")) _, ok, err = metering.Metering(context.TODO(), 1723487507) assert.NotNil(t, err) assert.False(t, ok) // metering success loader.EXPECT().LoadGCP(gomock.Any(), gomock.Any()).Return(billings, nil) store.EXPECT().Load(gomock.Any(), gomock.Any()).Return([]*models.BillingMetering{}, []*models.BillingReport{}, nil) store.EXPECT().Store(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil) store.EXPECT().Upload(gomock.Any(), gomock.Any()).Return("path", nil) path, ok, err := metering.Metering(context.TODO(), 1723487507) assert.Nil(t, err) assert.True(t, ok) assert.Equal(t, "path", path) } func TestEnsureUsedBy(t *testing.T) { // Test case 1: UserTagUsedBy is customer-tidb billing1 := &models.GCPBilling{ UserTagUsedBy: constants.UsedByCustomer, UserLabelUsedBy: "some-other-value", } result1 := ensureUsedBy(billing1) assert.Equal(t, constants.UsedByCustomer, result1) // Test case 2: UserTagUsedBy is not customer-tidb, fallback to UserLabelUsedBy billing2 := &models.GCPBilling{ UserTagUsedBy: "not-customer", UserLabelUsedBy: constants.UsedByCustomer, } result2 := ensureUsedBy(billing2) assert.Equal(t, constants.UsedByCustomer, result2) // Test case 3: UserTagUsedBy is empty, fallback to UserLabelUsedBy billing3 := &models.GCPBilling{ UserTagUsedBy: "", UserLabelUsedBy: constants.UsedByCustomer, } result3 := ensureUsedBy(billing3) assert.Equal(t, constants.UsedByCustomer, result3) // Test case 4: Both are empty billing4 := &models.GCPBilling{ UserTagUsedBy: "", UserLabelUsedBy: "", } result4 := ensureUsedBy(billing4) assert.Equal(t, "", result4) // Test case 5: UserTagUsedBy has some value but not customer-tidb billing5 := &models.GCPBilling{ UserTagUsedBy: "internal-use", UserLabelUsedBy: "fallback-value", } result5 := ensureUsedBy(billing5) assert.Equal(t, "fallback-value", result5) } func TestDealBillingAwsPremiumReserved(t *testing.T) { log, _ := zap.NewDevelopment() network := &NetworkMetering{ conf: &config.ServiceConfig{ Vendor: constants.VendorAWS, Env: constants.EnvDev, }, log: log, } billings := []*models.AWSBilling{ { LineItemOperation: "InterZone-In", ProductRegion: "us-west-2", LineItemProductCode: "AmazonEC2", ProductServicename: "AWS Data Transfer", ResourceTagsTenant: "456", ResourceTagsProject: "123", ResourceTagsCluster: "10663489082666642842", ResourceTagsUsedby: "", LineItemResourceId: "res-a", LineItemUnblendedCost: 2.1, MeteringServicePlan: constants.ServicePlanPremiumReserved, MeteringPremiumResMatchedOnly: true, }, { LineItemOperation: "InterZone-Out", ProductRegion: "us-west-2", LineItemProductCode: "AmazonEC2", ProductServicename: "AWS Data Transfer", ResourceTagsTenant: "456", ResourceTagsProject: "123", ResourceTagsCluster: "not-match", ResourceTagsSharedPool: "pool-1", ResourceTagsUsedby: "", LineItemResourceId: "res-b", LineItemUnblendedCost: 1.2, MeteringServicePlan: constants.ServicePlanPremiumReserved, MeteringPremiumResMatchedOnly: true, }, { LineItemOperation: "InterZone-Out", ProductRegion: "us-west-2", LineItemProductCode: "AmazonEC2", ProductServicename: "AWS Data Transfer", ResourceTagsTenant: "456", ResourceTagsProject: "123", ResourceTagsCluster: "not-match", ResourceTagsUsedby: "", LineItemResourceId: "res-c", LineItemUnblendedCost: 5.0, MeteringServicePlan: constants.ServicePlanPremiumReserved, MeteringPremiumResMatchedOnly: false, }, } res, err := network.DealBillingAWS(billings, nil, 1723487507) assert.NoError(t, err) assert.Len(t, res, 2) for _, r := range res { assert.Equal(t, constants.UsedByCustomer, r.UsedBy) } }