package nextgenmetering import ( "encoding/json" "strconv" "strings" "testing" "time" "github.com/pingcap/metering_sdk/common" "github.com/prometheus/client_golang/prometheus/testutil" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "github.com/tidbcloud/metering/pkg/common/config" "github.com/tidbcloud/metering/pkg/common/constants" "github.com/tidbcloud/metering/pkg/common/metrics" "github.com/tidbcloud/metering/pkg/common/models" "go.uber.org/zap" ) func newTestServiceConfig() *config.ServiceConfig { return &config.ServiceConfig{ Vendor: "aws", NextGenMeteringConfig: &config.NextGenMeteringConfig{ FeatureGateFilter: make(map[string]config.FeatureGateFilter), }, } } func newTestSKUGenerator(t *testing.T) *SKUGenerator { t.Helper() logger := zap.NewNop() return NewSKUGenerator(newTestServiceConfig(), logger) } func TestSKUGeneratorBasic(t *testing.T) { generator := newTestSKUGenerator(t) // Test RCU SKU generation for OLAP and OLTP t.Run("GenerateRCUSKU", func(t *testing.T) { enrichmentData := map[string]interface{}{ "minRcu": "5000", // intervalMin "maxRcu": "20000", // intervalMax } // Test combined RCU SKU with both OLAP and OLTP allRecord, err := generator.GenerateRCUSKU( "10000000", // clusterID "shared123", // sharedPoolID 6000, // tenantID 1700000, // timestamp "aws", // cloudProvider "us-west-2", // region "premium", // plan "pd", // component 1500, // totalRU (500 OLAP + 1000 OLTP) "RU", // ruUnit 500, // olapRU 1000, // oltpRU enrichmentData, time.Unix(1700000, 0), ) require.NoError(t, err) assert.Equal(t, models.NextGenSKUTypePublic, allRecord.SKUType) assert.Equal(t, "rcu:aws:us-west-2:premium:all:total", allRecord.Unit) assert.Equal(t, uint64(6000), allRecord.TenantID) assert.Equal(t, int64(1700000), allRecord.CostTimestamp) assert.Equal(t, int64(1), allRecord.Cost) // Cost represents quantity, should be 1 assert.Equal(t, "1", allRecord.UnitValue) // UnitValue is now fixed to "1" // Verify labels content matches JSON format assert.Contains(t, allRecord.Labels, "\"sourceUnit\":\"ru\"") assert.Contains(t, allRecord.Labels, "\"clusterID\":\"10000000\"") assert.Contains(t, allRecord.Labels, "\"sharedPoolID\":\"shared123\"") assert.Contains(t, allRecord.Labels, "\"intervalMin\":\"5000\"") assert.Contains(t, allRecord.Labels, "\"intervalMax\":\"20000\"") assert.Contains(t, allRecord.Labels, "\"intervalValue\":\"1500\"") // Total value assert.Contains(t, allRecord.Labels, "\"olapRU\":\"500\"") // OLAP value in labels assert.Contains(t, allRecord.Labels, "\"oltpRU\":\"1000\"") // OLTP value in labels // Test with only OLAP RU olapOnlyRecord, err := generator.GenerateRCUSKU( "10000000", // clusterID "shared123", // sharedPoolID 6000, // tenantID 1700000, // timestamp "aws", // cloudProvider "us-west-2", // region "premium", // plan "pd", // component 1000, // totalRU (1000 OLAP + 0 OLTP) "RU", // ruUnit 1000, // olapRU 0, // oltpRU enrichmentData, time.Unix(1700000, 0), ) require.NoError(t, err) assert.Equal(t, models.NextGenSKUTypePublic, olapOnlyRecord.SKUType) assert.Equal(t, "rcu:aws:us-west-2:premium:all:total", olapOnlyRecord.Unit) assert.Contains(t, olapOnlyRecord.Labels, "\"intervalValue\":\"1000\"") assert.Contains(t, olapOnlyRecord.Labels, "\"olapRU\":\"1000\"") assert.Contains(t, olapOnlyRecord.Labels, "\"oltpRU\":\"0\"") // Now always includes oltpRU, even if 0 }) // Test Data Transfer Out SKU generation t.Run("GenerateDataTransferOutSKU", func(t *testing.T) { enrichmentData := map[string]interface{}{} record, err := generator.GenerateDataTransferOutSKU( "10000000", // clusterID "shared123", // sharedPoolID 6000, // tenantID 1700000, // timestamp "aws", // cloudProvider "us-west-2", // region "premium", // plan "proxy", // component "", // networktype (empty for old version compatibility) 1000, // bytes (1KB for testing) "bytes", // sourceUnit enrichmentData, time.Unix(1700000, 0), ) require.NoError(t, err) assert.Equal(t, models.NextGenSKUTypePublic, record.SKUType) assert.Equal(t, "data-transfer-out:aws:us-west-2:premium:proxy", record.Unit) assert.Equal(t, uint64(6000), record.TenantID) assert.Equal(t, int64(1700000), record.CostTimestamp) assert.Equal(t, int64(1), record.Cost) // Cost represents quantity, should be 1 assert.Equal(t, "1000", record.UnitValue) // UnitValue contains the actual bytes value assert.Contains(t, record.Labels, "\"sourceUnit\":\"bytes\"") assert.Contains(t, record.Labels, "\"clusterID\":\"10000000\"") assert.Contains(t, record.Labels, "\"sharedPoolID\":\"shared123\"") // Test with private network type privateRecord, err := generator.GenerateDataTransferOutSKU( "10000000", "shared123", 6000, 1700000, "aws", "us-west-2", "premium", "proxy", models.NextGenProxyNetworkTypePrivate, 2000, "bytes", enrichmentData, time.Unix(1700000, 0), ) require.NoError(t, err) assert.Equal(t, models.NextGenSKUTypePublic, privateRecord.SKUType) assert.Equal(t, "data-transfer-out:aws:us-west-2:premium:proxy:private", privateRecord.Unit) assert.Equal(t, "2000", privateRecord.UnitValue) // Test with public network type publicRecord, err := generator.GenerateDataTransferOutSKU( "10000000", "shared123", 6000, 1700000, "aws", "us-west-2", "premium", "proxy", models.NextGenProxyNetworkTypePublic, 3000, "bytes", enrichmentData, time.Unix(1700000, 0), ) require.NoError(t, err) assert.Equal(t, models.NextGenSKUTypePublic, publicRecord.SKUType) assert.Equal(t, "data-transfer-out:aws:us-west-2:premium:proxy:public", publicRecord.Unit) assert.Equal(t, "3000", publicRecord.UnitValue) }) // Test Data Transfer Cross AZ SKU generation t.Run("GenerateDataTransferCrossAZSKU", func(t *testing.T) { enrichmentData := map[string]interface{}{} // Test with proxy component proxyRecord, err := generator.GenerateDataTransferCrossAZSKU( "10000000", // clusterID "shared123", // sharedPoolID 6000, // tenantID 1700000, // timestamp "aws", // cloudProvider "us-west-2", // region "premium", // plan "proxy", // component 1000, // bytes (1KB for testing) "bytes", // sourceUnit enrichmentData, time.Unix(1700000, 0), ) require.NoError(t, err) assert.Equal(t, models.NextGenSKUTypePublic, proxyRecord.SKUType) assert.Equal(t, "data-transfer-cross-az:aws:us-west-2:premium:proxy", proxyRecord.Unit) assert.Equal(t, uint64(6000), proxyRecord.TenantID) assert.Equal(t, int64(1700000), proxyRecord.CostTimestamp) assert.Equal(t, int64(1), proxyRecord.Cost) // Cost represents quantity, should be 1 assert.Equal(t, "1000", proxyRecord.UnitValue) // UnitValue contains the actual bytes value assert.Contains(t, proxyRecord.Labels, "\"sourceUnit\":\"bytes\"") assert.Contains(t, proxyRecord.Labels, "\"clusterID\":\"10000000\"") assert.Contains(t, proxyRecord.Labels, "\"sharedPoolID\":\"shared123\"") // Test with pd component pdRecord, err := generator.GenerateDataTransferCrossAZSKU( "10000000", // clusterID "shared123", // sharedPoolID 6000, // tenantID 1700000, // timestamp "aws", // cloudProvider "us-west-2", // region "premium", // plan "pd", // component 1000, // bytes (1KB for testing) "bytes", // sourceUnit enrichmentData, time.Unix(1700000, 0), ) require.NoError(t, err) assert.Equal(t, models.NextGenSKUTypePublic, pdRecord.SKUType) assert.Equal(t, "data-transfer-cross-az:aws:us-west-2:premium:pd", pdRecord.Unit) assert.Equal(t, uint64(6000), pdRecord.TenantID) assert.Equal(t, int64(1700000), pdRecord.CostTimestamp) assert.Equal(t, int64(1), pdRecord.Cost) // Cost represents quantity, should be 1 assert.Equal(t, "1000", pdRecord.UnitValue) // UnitValue contains the actual bytes value assert.Contains(t, pdRecord.Labels, "\"sourceUnit\":\"bytes\"") assert.Contains(t, pdRecord.Labels, "\"clusterID\":\"10000000\"") assert.Contains(t, pdRecord.Labels, "\"sharedPoolID\":\"shared123\"") }) // Test backup storage SKU generation for dfs-stats t.Run("GenerateBackupStorageSKUs", func(t *testing.T) { writtenBytes := uint64(2048) dataPoint := map[string]interface{}{ models.NextGenMeteringSourceName: "RfEngine", models.NextGenMeteringWrittenBytes: map[string]interface{}{ "value": writtenBytes, "unit": "Bytes", }, } records, err := generator.GenerateBackupStorageSKUs( dataPoint, "cluster-1", "shared-1", 6000, 1700000, "aws", "us-west-2", "premium", nil, time.Unix(1700000, 0), ) require.NoError(t, err) require.Len(t, records, 1) record := records[0] assert.Equal(t, "backup-storage:aws:us-west-2:premium:auto", record.Unit) assert.Equal(t, "2048", record.UnitValue) assert.Equal(t, models.NextGenSKUTypePublic, record.SKUType) var labels map[string]string require.NoError(t, json.Unmarshal([]byte(record.Labels), &labels)) assert.Equal(t, "bytes", labels[models.NextGenSKULabelKeySourceUnit]) assert.Equal(t, "auto", labels[models.NextGenSKULabelKeyBackupType]) assert.Equal(t, "rfengine", labels[models.NextGenSKULabelKeySourceName]) assert.Equal(t, "cluster-1", labels[models.NextGenSKULabelKeyClusterID]) assert.Equal(t, "shared-1", labels[models.NextGenSKULabelKeySharedPoolID]) }) t.Run("GenerateBackupStorageSKUs_NoWrittenBytes", func(t *testing.T) { dataPoint := map[string]interface{}{ models.NextGenMeteringSourceName: "rfengine", models.NextGenMeteringWrittenBytes: map[string]interface{}{ "value": 0, "unit": "Bytes", }, } records, err := generator.GenerateBackupStorageSKUs( dataPoint, "cluster-1", "shared-1", 6000, 1700000, "aws", "us-west-2", "premium", nil, time.Unix(1700000, 0), ) require.NoError(t, err) assert.Empty(t, records) }) // Test storage size SKU generation t.Run("GenerateStorageSizeSKUs", func(t *testing.T) { const gib = uint64(1024 * 1024 * 1024) rowBytes := uint64(2) * gib columnBytes := gib / 2 dataPoint := map[string]interface{}{ models.NextGenMeteringRowStorageSize: map[string]interface{}{ "value": rowBytes, "unit": "Bytes", }, models.NextGenMeteringColumnStorageSize: map[string]interface{}{ "value": columnBytes, "unit": "Bytes", }, models.NextGenMeteringSourceName: "pd", models.NextGenMeteringVersion: "1", } enrichmentData := map[string]interface{}{ models.NextGenMetaDiskType: "gp3", models.NextGenMetaColumnDataDiskType: "io2", models.NextGenMetaDiskIOPS: "8000", models.NextGenMetaDiskThroughput: "400", } records, err := generator.GenerateStorageSizeSKUs( dataPoint, "cluster-1", "", 6000, 1700000, "aws", "us-east-1", "premium", enrichmentData, time.Unix(1700000, 0), ) require.NoError(t, err) require.Len(t, records, 2) findRecord := func(unit string) *models.BillingReportV2Deltalake { for i := range records { if records[i].Unit == unit { return &records[i] } } return nil } rowStorageUnit := "disk-storage-size:aws:us-east-1:premium:row:standard:data:gp3" columnStorageUnit := "disk-storage-size:aws:us-east-1:premium:column:plus:data:io2" rowRecord := findRecord(rowStorageUnit) columnRecord := findRecord(columnStorageUnit) require.NotNil(t, rowRecord) require.NotNil(t, columnRecord) assert.Equal(t, strconv.FormatUint(rowBytes, 10), rowRecord.UnitValue) assert.Equal(t, int64(1), rowRecord.Cost) assert.Equal(t, models.NextGenSKUTypePublic, rowRecord.SKUType) var rowLabels map[string]string require.NoError(t, json.Unmarshal([]byte(rowRecord.Labels), &rowLabels)) assert.Equal(t, "pd", rowLabels[models.NextGenSKULabelKeySourceName]) assert.Equal(t, "bytes", rowLabels[models.NextGenSKULabelKeySourceUnit]) assert.Equal(t, "cluster-1", rowLabels[models.NextGenSKULabelKeyClusterID]) assert.Len(t, rowLabels, 3) testAbsent := func(labels map[string]string) { _, hasDiskType := labels[models.NextGenSKULabelKeyDiskType] assert.False(t, hasDiskType) _, hasStorageBytes := labels[models.NextGenSKULabelKeyStorageBytes] assert.False(t, hasStorageBytes) _, hasStorageCategory := labels[models.NextGenSKULabelKeyStorageCategory] assert.False(t, hasStorageCategory) _, hasStorageTier := labels[models.NextGenSKULabelKeyStorageTier] assert.False(t, hasStorageTier) _, hasStorageUsage := labels[models.NextGenSKULabelKeyStorageUsage] assert.False(t, hasStorageUsage) _, hasVersion := labels[models.NextGenSKULabelKeyVersion] assert.False(t, hasVersion) } testAbsent(rowLabels) assert.Equal(t, strconv.FormatUint(columnBytes, 10), columnRecord.UnitValue) assert.Equal(t, int64(1), columnRecord.Cost) assert.Equal(t, models.NextGenSKUTypePublic, columnRecord.SKUType) var columnLabels map[string]string require.NoError(t, json.Unmarshal([]byte(columnRecord.Labels), &columnLabels)) assert.Equal(t, "pd", columnLabels[models.NextGenSKULabelKeySourceName]) assert.Equal(t, "bytes", columnLabels[models.NextGenSKULabelKeySourceUnit]) assert.Equal(t, "cluster-1", columnLabels[models.NextGenSKULabelKeyClusterID]) assert.Len(t, columnLabels, 3) testAbsent(columnLabels) for _, record := range records { assert.NotContains(t, record.Unit, models.NextGenSKUDiskIOPS) } // var columnThroughputLabels map[string]string // require.NoError(t, json.Unmarshal([]byte(columnThroughputRecord.Labels), &columnThroughputLabels)) // assert.Equal(t, "pd", columnThroughputLabels[models.NextGenSKULabelKeySourceName]) // assert.Equal(t, "MBps", columnThroughputLabels[models.NextGenSKULabelKeySourceUnit]) // assert.Equal(t, "cluster-1", columnThroughputLabels[models.NextGenSKULabelKeyClusterID]) // assert.Len(t, columnThroughputLabels, 3) // testAbsent(columnThroughputLabels) }) t.Run("GenerateStorageSizeSKUsWithCMEKMetadataButGateDisabled", func(t *testing.T) { const gib = uint64(1024 * 1024 * 1024) rowBytes := uint64(2) * gib columnBytes := gib dataPoint := map[string]interface{}{ models.NextGenMeteringRowStorageSize: map[string]interface{}{ "value": rowBytes, "unit": "Bytes", }, models.NextGenMeteringColumnStorageSize: map[string]interface{}{ "value": columnBytes, "unit": "Bytes", }, models.NextGenMeteringSourceName: "pd", models.NextGenMeteringVersion: "1", } enrichmentData := map[string]interface{}{ models.NextGenMetaDiskType: "gp3", models.NextGenMetaColumnDataDiskType: "gp3", models.NextGenMetaStorageEncryptionEnabled: true, } records, err := generator.GenerateStorageSizeSKUs( dataPoint, "cluster-cmek", "", 6000, 1700000, "aws", "ap-northeast-1", "premium", enrichmentData, time.Unix(1700000, 0), ) require.NoError(t, err) require.Len(t, records, 2) expectedUnits := map[string]string{ "disk-storage-size:aws:ap-northeast-1:premium:row:standard:data:gp3": strconv.FormatUint(rowBytes, 10), "disk-storage-size:aws:ap-northeast-1:premium:column:standard:data:gp3": strconv.FormatUint(columnBytes, 10), } for _, record := range records { expectedValue, ok := expectedUnits[record.Unit] require.Truef(t, ok, "unexpected unit %s", record.Unit) assert.Equal(t, expectedValue, record.UnitValue) assert.Equal(t, int64(1), record.Cost) assert.Equal(t, models.NextGenSKUTypePublic, record.SKUType) var labels map[string]string require.NoError(t, json.Unmarshal([]byte(record.Labels), &labels)) assert.Equal(t, "pd", labels[models.NextGenSKULabelKeySourceName]) assert.Equal(t, "bytes", labels[models.NextGenSKULabelKeySourceUnit]) assert.Equal(t, "cluster-cmek", labels[models.NextGenSKULabelKeyClusterID]) } }) t.Run("GenerateStorageSizeSKUsWithCMEKGateEnabled", func(t *testing.T) { const gib = uint64(1024 * 1024 * 1024) const timestamp int64 = 1700000 rowBytes := uint64(2) * gib columnBytes := gib gateConf := newTestServiceConfig() gateConf.NextGenMeteringConfig.FeatureGateFilter[constants.FeatureGateNextGenCMEKStorageSKU] = config.FeatureGateFilter{ Clusters: []config.ClusterFilter{{ ClusterId: "cluster-cmek", SwitchTime: timestamp - 60, }}, } gateGenerator := NewSKUGenerator(gateConf, zap.NewNop()) dataPoint := map[string]interface{}{ models.NextGenMeteringRowStorageSize: map[string]interface{}{ "value": rowBytes, "unit": "Bytes", }, models.NextGenMeteringColumnStorageSize: map[string]interface{}{ "value": columnBytes, "unit": "Bytes", }, models.NextGenMeteringSourceName: "pd", models.NextGenMeteringVersion: "1", } enrichmentData := map[string]interface{}{ models.NextGenMetaDiskType: "gp3", models.NextGenMetaColumnDataDiskType: "gp3", models.NextGenMetaStorageEncryptionEnabled: true, models.NextGenMetaTenantID: "6000", } records, err := gateGenerator.GenerateStorageSizeSKUs( dataPoint, "cluster-cmek", "", 6000, timestamp, "aws", "ap-northeast-1", "premium", enrichmentData, time.Unix(timestamp, 0), ) require.NoError(t, err) require.Len(t, records, 2) expectedUnits := map[string]string{ "disk-storage-size-dual-encryption:aws:ap-northeast-1:premium:row:standard:data:gp3": strconv.FormatUint(rowBytes, 10), "disk-storage-size:aws:ap-northeast-1:premium:column:standard:data:gp3": strconv.FormatUint(columnBytes, 10), } for _, record := range records { expectedValue, ok := expectedUnits[record.Unit] require.Truef(t, ok, "unexpected unit %s", record.Unit) assert.Equal(t, expectedValue, record.UnitValue) assert.Equal(t, int64(1), record.Cost) assert.Equal(t, models.NextGenSKUTypePublic, record.SKUType) var labels map[string]string require.NoError(t, json.Unmarshal([]byte(record.Labels), &labels)) assert.Equal(t, "pd", labels[models.NextGenSKULabelKeySourceName]) assert.Equal(t, "bytes", labels[models.NextGenSKULabelKeySourceUnit]) assert.Equal(t, "cluster-cmek", labels[models.NextGenSKULabelKeyClusterID]) } }) t.Run("GenerateControlPlaneDiskSKUs", func(t *testing.T) { timestamp := int64(1700000) now := time.Unix(timestamp, 0) dataPoint := map[string]interface{}{ models.NextGenMeteringSourceName: "control-planel", models.NextGenMeteringTiKVDataDiskCount: map[string]interface{}{ "value": uint64(3), "unit": "count", }, models.NextGenMeteringTiKVDataDiskIOPS: map[string]interface{}{ "value": uint64(8000), "unit": "iops", }, models.NextGenMeteringTiKVDataDiskTP: map[string]interface{}{ "value": uint64(400), "unit": "mb/s", }, models.NextGenMeteringTiKVDataDiskType: "gp3", models.NextGenMeteringTiKVLogDiskCount: map[string]interface{}{ "value": uint64(3), "unit": "count", }, models.NextGenMeteringTiKVLogDiskIOPS: map[string]interface{}{ "value": uint64(3000), "unit": "iops", }, models.NextGenMeteringTiKVLogDiskTP: map[string]interface{}{ "value": uint64(125), "unit": "mb/s", }, models.NextGenMeteringTiKVLogDiskType: "gp3", models.NextGenMeteringTiFlashComputeDiskCount: map[string]interface{}{ "value": uint64(1), "unit": "count", }, models.NextGenMeteringTiFlashComputeDiskIOPS: map[string]interface{}{ "value": uint64(6000), "unit": "iops", }, models.NextGenMeteringTiFlashComputeDiskTP: map[string]interface{}{ "value": uint64(625), "unit": "mb/s", }, models.NextGenMeteringTiFlashComputeDiskType: "gp3", models.NextGenMeteringTiFlashWriteDiskCount: map[string]interface{}{ "value": uint64(1), "unit": "count", }, models.NextGenMeteringTiFlashWriteDiskIOPS: map[string]interface{}{ "value": uint64(6000), "unit": "iops", }, models.NextGenMeteringTiFlashWriteDiskTP: map[string]interface{}{ "value": uint64(625), "unit": "mb/s", }, models.NextGenMeteringTiFlashWriteDiskType: "gp3", } enrichmentData := map[string]interface{}{ models.NextGenMetaDiskType: "gp3", } records, err := generator.GenerateControlPlaneDiskSKUs( dataPoint, "cluster-1", "shared-1", 6000, timestamp, "aws", "us-west-2", "premium", enrichmentData, now, ) require.NoError(t, err) require.Len(t, records, 6) expected := map[string]struct { unitValue string cost int64 sourceUnit string diskType string }{ "disk-iops:aws:us-west-2:premium:row:standard:data:gp3": {unitValue: "8000", cost: 3, sourceUnit: "iops", diskType: "gp3"}, "disk-throughput:aws:us-west-2:premium:row:standard:data:gp3": {unitValue: "400", cost: 3, sourceUnit: "mb/s", diskType: "gp3"}, "disk-iops:aws:us-west-2:premium:row:standard:raft:gp3": {unitValue: "3000", cost: 3, sourceUnit: "iops", diskType: "gp3"}, "disk-throughput:aws:us-west-2:premium:row:standard:raft:gp3": {unitValue: "125", cost: 3, sourceUnit: "mb/s", diskType: "gp3"}, "disk-iops:aws:us-west-2:premium:column:standard:data:gp3": {unitValue: "6000", cost: 2, sourceUnit: "iops", diskType: "gp3"}, "disk-throughput:aws:us-west-2:premium:column:standard:data:gp3": {unitValue: "625", cost: 2, sourceUnit: "mb/s", diskType: "gp3"}, } seen := make(map[string]bool, len(records)) for _, record := range records { exp, ok := expected[record.Unit] assert.True(t, ok, "unexpected unit: %s", record.Unit) if !ok { continue } assert.Equal(t, exp.unitValue, record.UnitValue) assert.Equal(t, exp.cost, record.Cost) // Throughput SKUs should be private, others should be public if strings.Contains(record.Unit, "disk-throughput") { assert.Equal(t, models.NextGenSKUTypePrivate, record.SKUType) } else { assert.Equal(t, models.NextGenSKUTypePublic, record.SKUType) } var labels map[string]string require.NoError(t, json.Unmarshal([]byte(record.Labels), &labels)) assert.Equal(t, "control-planel", labels[models.NextGenSKULabelKeySourceName]) assert.Equal(t, exp.sourceUnit, labels[models.NextGenSKULabelKeySourceUnit]) assert.Equal(t, "cluster-1", labels[models.NextGenSKULabelKeyClusterID]) assert.Equal(t, "shared-1", labels[models.NextGenSKULabelKeySharedPoolID]) assert.Equal(t, exp.diskType, labels[models.NextGenSKULabelKeyDiskType]) if strings.Contains(record.Unit, "tiflash") { assert.Equal(t, "gp3", labels[models.NextGenSKULabelKeyTiFlashWriteDiskType]) _, hasWriteIOPS := labels[models.NextGenSKULabelKeyTiFlashWriteIOPS] _, hasWriteIOPSUnit := labels[models.NextGenSKULabelKeyTiFlashWriteIOPSUnit] _, hasWriteThroughput := labels[models.NextGenSKULabelKeyTiFlashWriteThroughput] _, hasWriteThroughputUnit := labels[models.NextGenSKULabelKeyTiFlashWriteThroughputUnit] if strings.Contains(record.Unit, models.NextGenSKUDiskIOPS) { assert.True(t, hasWriteIOPS) assert.True(t, hasWriteIOPSUnit) assert.False(t, hasWriteThroughput) assert.False(t, hasWriteThroughputUnit) } else if strings.Contains(record.Unit, models.NextGenSKUDiskThroughput) { assert.True(t, hasWriteThroughput) assert.True(t, hasWriteThroughputUnit) assert.False(t, hasWriteIOPS) assert.False(t, hasWriteIOPSUnit) } } seen[record.Unit] = true } assert.Len(t, seen, len(expected)) for unit := range expected { assert.True(t, seen[unit], "missing expected unit: %s", unit) } }) // Test edge case: empty clusterID should not include clusterID in labels t.Run("GenerateRCUSKU_EmptyClusterID", func(t *testing.T) { enrichmentData := map[string]interface{}{} record, err := generator.GenerateRCUSKU( "", // empty clusterID "shared123", // sharedPoolID 6000, // tenantID 1700000, // timestamp "aws", // cloudProvider "us-west-2", // region "premium", // plan "pd", // component 500, // totalRU "RU", // ruUnit 500, // olapRU 0, // oltpRU enrichmentData, time.Unix(1700000, 0), ) require.NoError(t, err) assert.Equal(t, models.NextGenSKUTypePrivate, record.SKUType) assert.Equal(t, "rcu:aws:us-west-2:premium:all:total", record.Unit) assert.Equal(t, uint64(6000), record.TenantID) assert.Contains(t, record.Labels, "\"sourceUnit\":\"ru\"") assert.Contains(t, record.Labels, "\"sharedPoolID\":\"shared123\"") assert.NotContains(t, record.Labels, "\"clusterID\"") // should not contain clusterID }) // Test edge case: system clusterID should set SKUType to private t.Run("GenerateRCUSKU_SystemClusterID", func(t *testing.T) { enrichmentData := map[string]interface{}{} record, err := generator.GenerateRCUSKU( models.NextGenMeteringSystemID, // system clusterID "SYSTEM" "shared123", // sharedPoolID 6000, // tenantID 1700000, // timestamp "aws", // cloudProvider "us-west-2", // region "premium", // plan "pd", // component 500, // totalRU "RU", // ruUnit 500, // olapRU 0, // oltpRU enrichmentData, time.Unix(1700000, 0), ) require.NoError(t, err) assert.Equal(t, models.NextGenSKUTypePrivate, record.SKUType) // Should be private for system ID assert.Equal(t, "rcu:aws:us-west-2:premium:all:total", record.Unit) assert.Equal(t, uint64(6000), record.TenantID) assert.Contains(t, record.Labels, "\"sourceUnit\":\"ru\"") assert.Contains(t, record.Labels, "\"sharedPoolID\":\"shared123\"") assert.NotContains(t, record.Labels, "\"clusterID\"") // system ID should not be included in labels }) // Test feature gate controlled unit format for RCUs t.Run("GenerateRCUSKU_FeatureGateClusterUnit", func(t *testing.T) { clusterID := "cluster-gate" const tenantID uint64 = 6000 const timestamp int64 = 1700000 noGateGenerator := newTestSKUGenerator(t) recordNoGate, err := noGateGenerator.GenerateRCUSKU( clusterID, "shared-1", tenantID, timestamp, "aws", "us-west-2", "premium", "pd", 1500, "RU", 500, 1000, map[string]interface{}{}, time.Unix(timestamp, 0), ) require.NoError(t, err) assert.Equal(t, "rcu:aws:us-west-2:premium:all:total", recordNoGate.Unit) gateConf := newTestServiceConfig() gateConf.NextGenMeteringConfig.FeatureGateFilter[constants.FeatureGateNextGenRCUInstance] = config.FeatureGateFilter{ Clusters: []config.ClusterFilter{{ ClusterId: clusterID, SwitchTime: timestamp - 60, }}, } gateGenerator := NewSKUGenerator(gateConf, zap.NewNop()) recordGate, err := gateGenerator.GenerateRCUSKU( clusterID, "shared-1", tenantID, timestamp, "aws", "us-west-2", "premium", "pd", 1500, "RU", 500, 1000, map[string]interface{}{}, time.Unix(timestamp, 0), ) require.NoError(t, err) assert.Equal(t, "rcu:aws:us-west-2:premium:all:total:"+clusterID, recordGate.Unit) }) t.Run("GenerateRCUSKU_MeteringMaxRCUGateOn", func(t *testing.T) { clusterID := "cluster-max-rcu-v2" const tenantID uint64 = 6000 const timestamp int64 = 1700000 gateConf := newTestServiceConfig() gateConf.NextGenMeteringConfig.FeatureGateFilter[constants.FeatureGateNextGenRCUMaxV2] = config.FeatureGateFilter{ Clusters: []config.ClusterFilter{{ ClusterId: clusterID, SwitchTime: timestamp - 60, }}, } gateGenerator := NewSKUGenerator(gateConf, zap.NewNop()) enrichmentData := map[string]interface{}{ models.NextGenMetaMaxRCU: "20000", models.NextGenMetaMeteringMaxRCU: "12000", } record, err := gateGenerator.GenerateRCUSKU( clusterID, "shared-1", tenantID, timestamp, "aws", "us-west-2", "premium", "pd", 1500, "RU", 500, 1000, enrichmentData, time.Unix(timestamp, 0), ) require.NoError(t, err) var labels map[string]string err = json.Unmarshal([]byte(record.Labels), &labels) require.NoError(t, err) assert.Equal(t, "12000", labels["intervalMax"]) }) t.Run("GenerateRCUSKU_MeteringMaxRCUEmptyStringDoesNotFallback", func(t *testing.T) { clusterID := "cluster-max-rcu-empty" const tenantID uint64 = 6000 const timestamp int64 = 1700000 gateConf := newTestServiceConfig() gateConf.NextGenMeteringConfig.FeatureGateFilter[constants.FeatureGateNextGenRCUMaxV2] = config.FeatureGateFilter{ Clusters: []config.ClusterFilter{{ ClusterId: clusterID, SwitchTime: timestamp - 60, }}, } gateGenerator := NewSKUGenerator(gateConf, zap.NewNop()) enrichmentData := map[string]interface{}{ models.NextGenMetaMaxRCU: "20000", models.NextGenMetaMeteringMaxRCU: "", } counter := metrics.MetricsNextGenMeteringMaxRCUFallbackCount.WithLabelValues( "aws", models.NextGenCategoryResourceManager, constants.MetricsNextGenReasonMeteringMaxRCUMissing, ) before := testutil.ToFloat64(counter) record, err := gateGenerator.GenerateRCUSKU( clusterID, "shared-1", tenantID, timestamp, "aws", "us-west-2", "premium", "pd", 1500, "RU", 500, 1000, enrichmentData, time.Unix(timestamp, 0), ) require.NoError(t, err) var labels map[string]string err = json.Unmarshal([]byte(record.Labels), &labels) require.NoError(t, err) value, ok := labels["intervalMax"] require.True(t, ok) assert.Equal(t, "", value) after := testutil.ToFloat64(counter) assert.Equal(t, before, after) }) t.Run("GenerateRCUSKU_MeteringMaxRCUFallback", func(t *testing.T) { clusterID := "cluster-max-rcu-fallback" const tenantID uint64 = 6000 const timestamp int64 = 1700000 gateConf := newTestServiceConfig() gateConf.NextGenMeteringConfig.FeatureGateFilter[constants.FeatureGateNextGenRCUMaxV2] = config.FeatureGateFilter{ Clusters: []config.ClusterFilter{{ ClusterId: clusterID, SwitchTime: timestamp - 60, }}, } gateGenerator := NewSKUGenerator(gateConf, zap.NewNop()) enrichmentData := map[string]interface{}{ models.NextGenMetaMaxRCU: "20000", } counter := metrics.MetricsNextGenMeteringMaxRCUFallbackCount.WithLabelValues( "aws", models.NextGenCategoryResourceManager, constants.MetricsNextGenReasonMeteringMaxRCUMissing, ) before := testutil.ToFloat64(counter) record, err := gateGenerator.GenerateRCUSKU( clusterID, "shared-1", tenantID, timestamp, "aws", "us-west-2", "premium", "pd", 1500, "RU", 500, 1000, enrichmentData, time.Unix(timestamp, 0), ) require.NoError(t, err) var labels map[string]string err = json.Unmarshal([]byte(record.Labels), &labels) require.NoError(t, err) assert.Equal(t, "20000", labels["intervalMax"]) after := testutil.ToFloat64(counter) assert.Equal(t, before+1, after) }) t.Run("backup storage v2 SKU generation", func(t *testing.T) { dataPoint := map[string]interface{}{ models.NextGenMeteringSourceName: models.NextGenSourceNamePD, models.NextGenMeteringOLAPRU: map[string]interface{}{ "value": uint64(1200), "unit": "ru", }, models.NextGenMeteringIncrementalWriteByte: map[string]interface{}{ "value": uint64(27947368), "unit": "bytes", }, } enrichmentData := map[string]interface{}{} now := time.Unix(1700006, 0) backupMeta := map[string]interface{}{ models.NextGenMetaBackupStorageV2: "T1", } records, err := generator.GenerateResourceManagerSKUs( dataPoint, "cluster-007", "pool-007", 11000, 1700006, "aws", "ap-southeast-1", "premium", enrichmentData, backupMeta, now, ) require.NoError(t, err) require.Len(t, records, 2) var backupRecord *models.BillingReportV2Deltalake for i := range records { if strings.HasPrefix(records[i].Unit, models.NextGenSKUBackupStorage+":") { backupRecord = &records[i] break } } require.NotNil(t, backupRecord) assert.Equal(t, "backup-storage:aws:ap-southeast-1:premium:T1", backupRecord.Unit) assert.Equal(t, "27947368", backupRecord.UnitValue) var labels map[string]string require.NoError(t, json.Unmarshal([]byte(backupRecord.Labels), &labels)) assert.Equal(t, "bytes", labels[models.NextGenSKULabelKeySourceUnit]) assert.Equal(t, "cluster-007", labels[models.NextGenSKULabelKeyClusterID]) }) t.Run("backup storage v2 SKU skipped when extMeta is daily", func(t *testing.T) { dataPoint := map[string]interface{}{ models.NextGenMeteringSourceName: models.NextGenSourceNamePD, models.NextGenMeteringOLAPRU: map[string]interface{}{ "value": uint64(1200), "unit": "ru", }, models.NextGenMeteringIncrementalWriteByte: map[string]interface{}{ "value": uint64(27947368), "unit": "bytes", }, } backupMeta := map[string]interface{}{ models.NextGenMetaBackupStorageV2: models.NextGenBackupTypeFlex, } records, err := generator.GenerateResourceManagerSKUs( dataPoint, "cluster-007", "pool-007", 11000, 1700006, "aws", "ap-southeast-1", "premium", map[string]interface{}{}, backupMeta, time.Unix(1700006, 0), ) require.NoError(t, err) for _, r := range records { assert.False(t, strings.HasPrefix(r.Unit, models.NextGenSKUBackupStorage+":"), "backup-storage record should be absent when extMeta is daily (handled by flex-backup flow)") } }) } func TestGenerateManualBackupStorageSKUs(t *testing.T) { generator := newTestSKUGenerator(t) now := time.Unix(1700000, 0) t.Run("generates manual backup storage sku in gb", func(t *testing.T) { dataPoint := map[string]interface{}{ models.NextGenMeteringLogicalClusterID: "cluster-1", models.NextGenMeteringBackupID: "backup-1", models.NextGenMeteringBackupSizeGB: map[string]interface{}{ "value": uint64(128), "unit": "GB", }, } records, err := generator.GenerateManualBackupStorageSKUs( dataPoint, "cluster-1", "shared-1", 6000, 1700000, "aws", "ap-south-1", "premium", map[string]interface{}{}, now, ) require.NoError(t, err) require.Len(t, records, 1) record := records[0] assert.Equal(t, "backup-storage:aws:ap-south-1:premium:manual", record.Unit) assert.Equal(t, "128", record.UnitValue) assert.Equal(t, int64(1), record.Cost) assert.Equal(t, models.NextGenSKUTypePublic, record.SKUType) assert.Equal(t, uint64(6000), record.TenantID) assert.Equal(t, int64(1700000), record.CostTimestamp) var labels map[string]string require.NoError(t, json.Unmarshal([]byte(record.Labels), &labels)) assert.Equal(t, "gb", labels[models.NextGenSKULabelKeySourceUnit]) assert.Equal(t, models.NextGenBackupTypeManual, labels[models.NextGenSKULabelKeyBackupType]) assert.Equal(t, "backup-1", labels[models.NextGenSKULabelKeyBackupID]) assert.Equal(t, "cluster-1", labels[models.NextGenSKULabelKeyClusterID]) assert.Equal(t, "shared-1", labels[models.NextGenSKULabelKeySharedPoolID]) }) t.Run("skips zero backup size", func(t *testing.T) { dataPoint := map[string]interface{}{ models.NextGenMeteringBackupID: "backup-zero", models.NextGenMeteringBackupSizeGB: map[string]interface{}{ "value": uint64(0), "unit": "GB", }, } records, err := generator.GenerateManualBackupStorageSKUs( dataPoint, "cluster-1", "shared-1", 6000, 1700000, "aws", "ap-south-1", "premium", map[string]interface{}{}, now, ) require.NoError(t, err) assert.Empty(t, records) }) t.Run("skips missing backup size", func(t *testing.T) { dataPoint := map[string]interface{}{ models.NextGenMeteringBackupID: "backup-missing-size", } records, err := generator.GenerateManualBackupStorageSKUs( dataPoint, "cluster-1", "shared-1", 6000, 1700000, "aws", "ap-south-1", "premium", map[string]interface{}{}, now, ) require.NoError(t, err) assert.Empty(t, records) }) t.Run("skips malformed backup size", func(t *testing.T) { dataPoint := map[string]interface{}{ models.NextGenMeteringBackupID: "backup-bad-size", models.NextGenMeteringBackupSizeGB: "128", } records, err := generator.GenerateManualBackupStorageSKUs( dataPoint, "cluster-1", "shared-1", 6000, 1700000, "aws", "ap-south-1", "premium", map[string]interface{}{}, now, ) require.NoError(t, err) assert.Empty(t, records) }) t.Run("system cluster emits private sku without cluster label", func(t *testing.T) { dataPoint := map[string]interface{}{ models.NextGenMeteringBackupSizeGB: map[string]interface{}{ "value": uint64(64), "unit": "GB", }, } records, err := generator.GenerateManualBackupStorageSKUs( dataPoint, models.NextGenMeteringSystemID, "", 6000, 1700000, "aws", "ap-south-1", "premium", map[string]interface{}{}, now, ) require.NoError(t, err) require.Len(t, records, 1) record := records[0] assert.Equal(t, models.NextGenSKUTypePrivate, record.SKUType) var labels map[string]string require.NoError(t, json.Unmarshal([]byte(record.Labels), &labels)) _, hasClusterID := labels[models.NextGenSKULabelKeyClusterID] assert.False(t, hasClusterID) _, hasSharedPoolID := labels[models.NextGenSKULabelKeySharedPoolID] assert.False(t, hasSharedPoolID) _, hasBackupID := labels[models.NextGenSKULabelKeyBackupID] assert.False(t, hasBackupID) }) t.Run("empty cluster id emits private sku", func(t *testing.T) { dataPoint := map[string]interface{}{ models.NextGenMeteringBackupSizeGB: map[string]interface{}{ "value": uint64(32), "unit": "GB", }, } records, err := generator.GenerateManualBackupStorageSKUs( dataPoint, "", "shared-1", 6000, 1700000, "aws", "ap-south-1", "premium", map[string]interface{}{}, now, ) require.NoError(t, err) require.Len(t, records, 1) assert.Equal(t, models.NextGenSKUTypePrivate, records[0].SKUType) }) } func TestGenerateFlexBackupStorageSKUs(t *testing.T) { generator := newTestSKUGenerator(t) now := time.Unix(1700000, 0) t.Run("generates flex backup storage sku with daily suffix", func(t *testing.T) { dataPoint := map[string]interface{}{ models.NextGenMeteringLogicalClusterID: "cluster-1", models.NextGenMeteringBackupID: "backup-1", models.NextGenMeteringBackupSizeGB: map[string]interface{}{ "value": uint64(128), "unit": "GB", }, } records, err := generator.GenerateFlexBackupStorageSKUs( dataPoint, "cluster-1", "shared-1", 6000, 1700000, "aws", "ap-northeast-1", "premium", map[string]interface{}{}, now, ) require.NoError(t, err) require.Len(t, records, 1) record := records[0] assert.Equal(t, "backup-storage:aws:ap-northeast-1:premium:daily", record.Unit) assert.Equal(t, "128", record.UnitValue) assert.Equal(t, int64(1), record.Cost) assert.Equal(t, models.NextGenSKUTypePublic, record.SKUType) assert.Equal(t, uint64(6000), record.TenantID) assert.Equal(t, int64(1700000), record.CostTimestamp) var labels map[string]string require.NoError(t, json.Unmarshal([]byte(record.Labels), &labels)) assert.Equal(t, "gb", labels[models.NextGenSKULabelKeySourceUnit]) assert.Equal(t, models.NextGenBackupTypeFlex, labels[models.NextGenSKULabelKeyBackupType]) assert.Equal(t, "backup-1", labels[models.NextGenSKULabelKeyBackupID]) assert.Equal(t, "cluster-1", labels[models.NextGenSKULabelKeyClusterID]) assert.Equal(t, "shared-1", labels[models.NextGenSKULabelKeySharedPoolID]) }) t.Run("skips zero backup size", func(t *testing.T) { dataPoint := map[string]interface{}{ models.NextGenMeteringBackupID: "backup-zero", models.NextGenMeteringBackupSizeGB: map[string]interface{}{ "value": uint64(0), "unit": "GB", }, } records, err := generator.GenerateFlexBackupStorageSKUs( dataPoint, "cluster-1", "shared-1", 6000, 1700000, "aws", "ap-northeast-1", "premium", map[string]interface{}{}, now, ) require.NoError(t, err) assert.Empty(t, records) }) t.Run("skips missing backup size", func(t *testing.T) { dataPoint := map[string]interface{}{ models.NextGenMeteringBackupID: "backup-missing-size", } records, err := generator.GenerateFlexBackupStorageSKUs( dataPoint, "cluster-1", "shared-1", 6000, 1700000, "aws", "ap-northeast-1", "premium", map[string]interface{}{}, now, ) require.NoError(t, err) assert.Empty(t, records) }) t.Run("system cluster emits private sku without cluster label", func(t *testing.T) { dataPoint := map[string]interface{}{ models.NextGenMeteringBackupSizeGB: map[string]interface{}{ "value": uint64(64), "unit": "GB", }, } records, err := generator.GenerateFlexBackupStorageSKUs( dataPoint, models.NextGenMeteringSystemID, "", 6000, 1700000, "aws", "ap-northeast-1", "premium", map[string]interface{}{}, now, ) require.NoError(t, err) require.Len(t, records, 1) assert.Equal(t, models.NextGenSKUTypePrivate, records[0].SKUType) var labels map[string]string require.NoError(t, json.Unmarshal([]byte(records[0].Labels), &labels)) _, hasClusterID := labels[models.NextGenSKULabelKeyClusterID] assert.False(t, hasClusterID) _, hasSharedPoolID := labels[models.NextGenSKULabelKeySharedPoolID] assert.False(t, hasSharedPoolID) _, hasBackupID := labels[models.NextGenSKULabelKeyBackupID] assert.False(t, hasBackupID) }) } // TestRCUSKUJSONFormat tests the exact JSON format output matching the provided examples func TestRCUSKUJSONFormat(t *testing.T) { generator := newTestSKUGenerator(t) // Test data matching the provided JSON examples enrichmentData := map[string]interface{}{ "minRcu": "5000", // intervalMin "maxRcu": "20000", // intervalMax } // Test combined RCU SKU JSON format t.Run("Combined_RCU_JSON_Format", func(t *testing.T) { allRecord, err := generator.GenerateRCUSKU( "10000000", // clusterID "shared123", // sharedPoolID 6000, // tenantID 1700000, // timestamp "aws", // cloudProvider "us-west-2", // region "premium", // plan "pd", // component 2000, // totalRU (1000 OLAP + 1000 OLTP) "RU", // ruUnit 1000, // olapRU 1000, // oltpRU enrichmentData, time.Unix(1700000, 0), ) require.NoError(t, err) // Verify basic fields assert.Equal(t, uint64(6000), allRecord.TenantID) assert.Equal(t, int64(1), allRecord.Cost) // Cost represents quantity, should be 1 assert.Equal(t, "1", allRecord.UnitValue) // UnitValue is now fixed to "1" assert.Equal(t, int64(1700000), allRecord.CostTimestamp) assert.Equal(t, "rcu:aws:us-west-2:premium:all:total", allRecord.Unit) // Parse and verify labels JSON var labels map[string]string err = json.Unmarshal([]byte(allRecord.Labels), &labels) require.NoError(t, err) // Verify labels content completely matches JSON example assert.Equal(t, "10000000", labels["clusterID"]) assert.Equal(t, "5000", labels["intervalMin"]) assert.Equal(t, "20000", labels["intervalMax"]) assert.Equal(t, "ru", labels["sourceUnit"]) assert.Equal(t, "shared123", labels["sharedPoolID"]) assert.Equal(t, "2000", labels["intervalValue"]) // Total RCU value assert.Equal(t, "1000", labels["olapRU"]) // OLAP RU in labels assert.Equal(t, "1000", labels["oltpRU"]) // OLTP RU in labels }) // Test edge case: when sharedPoolID is empty, it should not be included in labels t.Run("Empty_SharedPoolID", func(t *testing.T) { record, err := generator.GenerateRCUSKU( "10000000", // clusterID "", // empty sharedPoolID 6000, // tenantID 1700000, // timestamp "aws", // cloudProvider "us-west-2", // region "premium", // plan "pd", // component 1000, // totalRU "RU", // ruUnit 1000, // olapRU 0, // oltpRU enrichmentData, time.Unix(1700000, 0), ) require.NoError(t, err) // Parse labels JSON var labels map[string]string err = json.Unmarshal([]byte(record.Labels), &labels) require.NoError(t, err) // Verify sharedPoolID is not included _, hasSharedPoolID := labels["sharedPoolID"] assert.False(t, hasSharedPoolID, "Labels should not contain sharedPoolID when it's empty") // But should contain other required fields assert.Equal(t, "10000000", labels["clusterID"]) assert.Equal(t, "5000", labels["intervalMin"]) assert.Equal(t, "20000", labels["intervalMax"]) assert.Equal(t, "ru", labels["sourceUnit"]) assert.Equal(t, "1000", labels["intervalValue"]) // Real value is in intervalValue assert.Equal(t, "1000", labels["olapRU"]) }) // Test edge case: when intervalMin/intervalMax do not exist, they should not be included in labels t.Run("Missing_Interval_Values", func(t *testing.T) { // Empty enrichmentData, does not contain minRcu/maxRcu emptyEnrichmentData := map[string]interface{}{} record, err := generator.GenerateRCUSKU( "10000000", // clusterID "shared123", // sharedPoolID 6000, // tenantID 1700000, // timestamp "aws", // cloudProvider "us-west-2", // region "premium", // plan "pd", // component 1000, // totalRU "RU", // ruUnit 1000, // olapRU 0, // oltpRU emptyEnrichmentData, time.Unix(1700000, 0), ) require.NoError(t, err) // Parse labels JSON var labels map[string]string err = json.Unmarshal([]byte(record.Labels), &labels) require.NoError(t, err) // Verify intervalMin is always 5000 (fixed value), but intervalMax should not be included when not provided assert.Equal(t, "5000", labels["intervalMin"], "intervalMin should always be 5000 (fixed)") _, hasIntervalMax := labels["intervalMax"] assert.False(t, hasIntervalMax, "Labels should not contain intervalMax when not provided") // But should contain other basic fields assert.Equal(t, "10000000", labels["clusterID"]) assert.Equal(t, "shared123", labels["sharedPoolID"]) assert.Equal(t, "ru", labels["sourceUnit"]) assert.Equal(t, "1000", labels["intervalValue"]) // Real value is in intervalValue assert.Equal(t, "1000", labels["olapRU"]) assert.Equal(t, "0", labels["oltpRU"]) // Now always includes oltpRU }) // Test that olapRU and oltpRU are always recorded, even when 0 t.Run("Zero_RU_Values_Always_Recorded", func(t *testing.T) { enrichmentData := map[string]interface{}{ "maxRcu": "10000", } // Test with both values as 0 record, err := generator.GenerateRCUSKU( "test-cluster", "", 1000, 1700000, "gcp", "us-central1", "basic", "pd", 0, // totalRU = 0 "RU", 0, // olapRU = 0 0, // oltpRU = 0 enrichmentData, time.Unix(1700000, 0), ) require.NoError(t, err) var labels map[string]string err = json.Unmarshal([]byte(record.Labels), &labels) require.NoError(t, err) // Both should be recorded as "0" assert.Equal(t, "0", labels["olapRU"], "olapRU should be recorded even when 0") assert.Equal(t, "0", labels["oltpRU"], "oltpRU should be recorded even when 0") assert.Equal(t, "0", labels["intervalValue"]) }) // Test that minRU is read from enrichmentData, with fallback to fixed 5000 t.Run("Dynamic_MinRU_Value", func(t *testing.T) { // Test 1: No minRcu in enrichmentData — should fall back to fixed value emptyEnrichment := map[string]interface{}{} record1, err := generator.GenerateRCUSKU( "cluster-1", "", 2000, 1700000, "aws", "us-east-1", "premium", "pd", 1000, "RU", 1000, 0, emptyEnrichment, time.Unix(1700000, 0), ) require.NoError(t, err) var labels1 map[string]string err = json.Unmarshal([]byte(record1.Labels), &labels1) require.NoError(t, err) assert.Equal(t, models.NextgenFixMinRU, labels1["intervalMin"], "intervalMin should fall back to fixed value when not in enrichmentData") // Test 2: minRcu from enrichmentData should be used differentEnrichment := map[string]interface{}{ "minRcu": "8000", "maxRcu": "20000", } record2, err := generator.GenerateRCUSKU( "cluster-2", "", 3000, 1700000, "azure", "eastus", "basic", "tikv", 2000, "RU", 1500, 500, differentEnrichment, time.Unix(1700000, 0), ) require.NoError(t, err) var labels2 map[string]string err = json.Unmarshal([]byte(record2.Labels), &labels2) require.NoError(t, err) assert.Equal(t, "8000", labels2["intervalMin"], "intervalMin should come from enrichmentData") assert.Equal(t, "20000", labels2["intervalMax"], "intervalMax should still come from enrichmentData") }) } // TestCompleteJSONStructure tests the complete JSON structure output func TestCompleteJSONStructure(t *testing.T) { generator := newTestSKUGenerator(t) enrichmentData := map[string]interface{}{ "minRcu": "5000", "maxRcu": "20000", } // Test complete JSON structure for combined RCU t.Run("Complete_Combined_RCU_Structure", func(t *testing.T) { record, err := generator.GenerateRCUSKU( "10000000", // clusterID "shared123", // sharedPoolID 6000, // tenantID 1700000, // timestamp "aws", // cloudProvider "us-west-2", // region "premium", // plan "pd", // component 2000, // totalRU (1000 OLAP + 1000 OLTP) "RU", // ruUnit 1000, // olapRU 1000, // oltpRU enrichmentData, time.Unix(1700000, 0), ) require.NoError(t, err) // Create expected structure to compare with actual output expectedJSON := map[string]interface{}{ "tenant_id": float64(6000), // JSON numbers are parsed as float64 "cost": float64(1), // Cost represents quantity, should be 1 "cost_time_stamp": float64(1700000), "unit": "rcu:aws:us-west-2:premium:all:total", // Changed to "all" type "unit_value": "1", // UnitValue is now fixed to "1" } expectedLabels := map[string]string{ "clusterID": "10000000", "intervalMin": "5000", "intervalMax": "20000", "intervalValue": "2000", // Total value: 1000 OLAP + 1000 OLTP = 2000 "sourceUnit": "ru", "sharedPoolID": "shared123", "olapRU": "1000", // OLAP RU in labels "oltpRU": "1000", // OLTP RU in labels } // Verify basic fields assert.Equal(t, uint64(expectedJSON["tenant_id"].(float64)), record.TenantID) assert.Equal(t, int64(expectedJSON["cost"].(float64)), record.Cost) assert.Equal(t, int64(expectedJSON["cost_time_stamp"].(float64)), record.CostTimestamp) assert.Equal(t, expectedJSON["unit"].(string), record.Unit) assert.Equal(t, expectedJSON["unit_value"].(string), record.UnitValue) // Verify labels var actualLabels map[string]string err = json.Unmarshal([]byte(record.Labels), &actualLabels) require.NoError(t, err) assert.Equal(t, expectedLabels, actualLabels) // Verify cost_id is not empty assert.NotEmpty(t, record.CostID) // Verify other required fields assert.Equal(t, "v2", record.SchemeVersion) assert.Equal(t, models.NextGenSKUTypePublic, record.SKUType) assert.NotZero(t, record.CreatedAt) assert.NotEmpty(t, record.HourDimension) t.Logf("Generated combined RCU record: tenant_id=%d, cost=%d, cost_time_stamp=%d, unit=%s", record.TenantID, record.Cost, record.CostTimestamp, record.Unit) t.Logf("Generated labels: %s", record.Labels) }) // Test complete JSON structure with only OLTP t.Run("Complete_OLTP_Only_Structure", func(t *testing.T) { record, err := generator.GenerateRCUSKU( "10000000", // clusterID "shared123", // sharedPoolID 6000, // tenantID 1700000, // timestamp "aws", // cloudProvider "us-west-2", // region "premium", // plan "pd", // component 1000, // totalRU "RU", // ruUnit 0, // olapRU 1000, // oltpRU enrichmentData, time.Unix(1700000, 0), ) require.NoError(t, err) // Verify unit format - should be "all" type assert.Equal(t, "rcu:aws:us-west-2:premium:all:total", record.Unit) // Verify other fields are same as OLAP assert.Equal(t, uint64(6000), record.TenantID) assert.Equal(t, int64(1), record.Cost) // Cost represents quantity, should be 1 assert.Equal(t, "1", record.UnitValue) // UnitValue is now fixed to "1" assert.Equal(t, int64(1700000), record.CostTimestamp) // Verify labels structure var actualLabels map[string]string err = json.Unmarshal([]byte(record.Labels), &actualLabels) require.NoError(t, err) expectedLabels := map[string]string{ "clusterID": "10000000", "intervalMin": "5000", "intervalMax": "20000", "intervalValue": "1000", // Real value is in intervalValue "sourceUnit": "ru", "sharedPoolID": "shared123", "olapRU": "0", // OLAP RU always recorded, even if 0 "oltpRU": "1000", // OLTP RU in labels } assert.Equal(t, expectedLabels, actualLabels) t.Logf("Generated OLTP record: tenant_id=%d, cost=%d, cost_time_stamp=%d, unit=%s", record.TenantID, record.Cost, record.CostTimestamp, record.Unit) t.Logf("Generated labels: %s", record.Labels) }) } // TestDXFSKUGeneration tests DXF SKU generation func TestDXFSKUGeneration(t *testing.T) { generator := newTestSKUGenerator(t) t.Run("GenerateDXFSKUs", func(t *testing.T) { // Create sample DXF data point dataPoint := map[string]interface{}{ "cluster_id": "10072959209283210116", "task_id": 30001, "task_type": "import-into", "version": "1", "source_name": "dxf", "get_external_requests": 8, "put_external_requests": 24, "read_data_bytes": map[string]interface{}{ "value": 109689961, "unit": "bytes", }, "write_data_bytes": map[string]interface{}{ "value": 37004046, "unit": "bytes", }, } enrichmentData := map[string]interface{}{} now := time.Unix(1758533100, 0) records, err := generator.GenerateDXFSKUs( dataPoint, "10072959209283210116", // clusterID "1041189107389498", // sharedPoolID 6000, // tenantID 1758533100, // timestamp "aws", // cloudProvider "us-west-2", // region "premium", // plan enrichmentData, now, ) require.NoError(t, err) require.Len(t, records, 1) record := records[0] // Verify basic fields assert.Equal(t, "task-ru:aws:us-west-2:premium:import-into", record.Unit) assert.Equal(t, uint64(6000), record.TenantID) assert.Equal(t, int64(1758533100), record.CostTimestamp) assert.Equal(t, int64(1), record.Cost) assert.Equal(t, "1", record.UnitValue) assert.Equal(t, models.NextGenSKUTypePrivate, record.SKUType) // DXF should be private // Parse and verify labels var labels map[string]string err = json.Unmarshal([]byte(record.Labels), &labels) require.NoError(t, err) // Verify all DXF fields are in labels assert.Equal(t, "10072959209283210116", labels["clusterID"]) assert.Equal(t, "1041189107389498", labels["sharedPoolID"]) assert.Equal(t, "30001", labels["taskID"]) assert.Equal(t, "import-into", labels["taskType"]) assert.Equal(t, "1", labels["version"]) assert.Equal(t, "8", labels["getExternalRequests"]) assert.Equal(t, "24", labels["putExternalRequests"]) assert.Equal(t, "109689961", labels["readDataBytes"]) assert.Equal(t, "37004046", labels["writeDataBytes"]) // Verify correct number of fields expectedFieldCount := 9 assert.Equal(t, expectedFieldCount, len(labels)) t.Logf("Generated DXF record: %+v", record) t.Logf("Generated labels: %s", record.Labels) }) t.Run("GenerateTaskRUSKU", func(t *testing.T) { dataPoint := map[string]interface{}{ "task_id": 12345, "task_type": "export-data", "version": "2", "get_external_requests": 15, "put_external_requests": 30, "read_data_bytes": map[string]interface{}{ "value": 1000000, "unit": "bytes", }, "write_data_bytes": map[string]interface{}{ "value": 500000, "unit": "bytes", }, } enrichmentData := map[string]interface{}{} now := time.Unix(1700000, 0) record, err := generator.GenerateTaskRUSKU( "cluster123", // clusterID "pool456", // sharedPoolID 7000, // tenantID 1700000, // timestamp "gcp", // cloudProvider "us-central1", // region "enterprise", // plan "dxf", // component "export-data", // taskType dataPoint, enrichmentData, now, ) require.NoError(t, err) // Verify basic structure assert.Equal(t, "task-ru:gcp:us-central1:enterprise:export-data", record.Unit) assert.Equal(t, uint64(7000), record.TenantID) assert.Equal(t, int64(1700000), record.CostTimestamp) assert.Equal(t, int64(1), record.Cost) assert.Equal(t, "1", record.UnitValue) assert.Equal(t, models.NextGenSKUTypePrivate, record.SKUType) // Parse labels var labels map[string]string err = json.Unmarshal([]byte(record.Labels), &labels) require.NoError(t, err) // Verify specific label values assert.Equal(t, "cluster123", labels["clusterID"]) assert.Equal(t, "pool456", labels["sharedPoolID"]) assert.Equal(t, "12345", labels["taskID"]) assert.Equal(t, "export-data", labels["taskType"]) assert.Equal(t, "2", labels["version"]) assert.Equal(t, "15", labels["getExternalRequests"]) assert.Equal(t, "30", labels["putExternalRequests"]) assert.Equal(t, "1000000", labels["readDataBytes"]) assert.Equal(t, "500000", labels["writeDataBytes"]) }) // Test edge cases t.Run("DXF_EmptyClusterID", func(t *testing.T) { dataPoint := map[string]interface{}{ "task_id": 999, "task_type": "test-task", } record, err := generator.GenerateTaskRUSKU( "", // empty clusterID "shared_pool_1", // sharedPoolID 8000, // tenantID 1700000, // timestamp "azure", // cloudProvider "eastus", // region "standard", // plan "dxf", // component "test-task", // taskType dataPoint, map[string]interface{}{}, time.Unix(1700000, 0), ) require.NoError(t, err) // Parse labels var labels map[string]string err = json.Unmarshal([]byte(record.Labels), &labels) require.NoError(t, err) // clusterID should not be present when empty _, hasClusterID := labels["clusterID"] assert.False(t, hasClusterID) // But sharedPoolID should be present assert.Equal(t, "shared_pool_1", labels["sharedPoolID"]) assert.Equal(t, "999", labels["taskID"]) assert.Equal(t, "test-task", labels["taskType"]) }) t.Run("DXF_MissingOptionalFields", func(t *testing.T) { // Minimal data point with only required fields dataPoint := map[string]interface{}{ "task_type": "minimal-task", } record, err := generator.GenerateTaskRUSKU( "cluster_min", "pool_min", 9000, 1700000, "aws", "us-east-1", "basic", "dxf", "minimal-task", dataPoint, map[string]interface{}{}, time.Unix(1700000, 0), ) require.NoError(t, err) // Parse labels var labels map[string]string err = json.Unmarshal([]byte(record.Labels), &labels) require.NoError(t, err) // Should contain basic fields assert.Equal(t, "cluster_min", labels["clusterID"]) assert.Equal(t, "pool_min", labels["sharedPoolID"]) assert.Equal(t, "minimal-task", labels["taskType"]) // Should not contain missing optional fields _, hasTaskID := labels["taskID"] assert.False(t, hasTaskID) _, hasVersion := labels["version"] assert.False(t, hasVersion) _, hasGetExtReq := labels["getExternalRequests"] assert.False(t, hasGetExtReq) }) } // TestConvertToMeteringValue tests the ConvertToMeteringValue function comprehensively func TestConvertToMeteringValue(t *testing.T) { generator := newTestSKUGenerator(t) tests := []struct { name string input interface{} expectedValue *common.MeteringValue expectedOk bool }{ { name: "nil input", input: nil, expectedValue: nil, expectedOk: false, }, { name: "already MeteringValue pointer", input: &common.MeteringValue{ Value: 100, Unit: "bytes", }, expectedValue: &common.MeteringValue{ Value: 100, Unit: "bytes", }, expectedOk: true, }, { name: "map with uint64 value", input: map[string]interface{}{ "value": uint64(500), "unit": "ru", }, expectedValue: &common.MeteringValue{ Value: 500, Unit: "ru", }, expectedOk: true, }, { name: "map with positive int64 value", input: map[string]interface{}{ "value": int64(1000), "unit": "bytes", }, expectedValue: &common.MeteringValue{ Value: 1000, Unit: "bytes", }, expectedOk: true, }, { name: "map with negative int64 value", input: map[string]interface{}{ "value": int64(-100), "unit": "bytes", }, expectedValue: &common.MeteringValue{ Value: 0, // negative values are not converted, leaving Value as 0 Unit: "bytes", }, expectedOk: true, }, { name: "map with positive int value", input: map[string]interface{}{ "value": int(750), "unit": "requests", }, expectedValue: &common.MeteringValue{ Value: 750, Unit: "requests", }, expectedOk: true, }, { name: "map with negative int value", input: map[string]interface{}{ "value": int(-50), "unit": "requests", }, expectedValue: &common.MeteringValue{ Value: 0, // negative values are not converted, leaving Value as 0 Unit: "requests", }, expectedOk: true, }, { name: "map with positive float64 value", input: map[string]interface{}{ "value": float64(123.456), "unit": "mb", }, expectedValue: &common.MeteringValue{ Value: 123, Unit: "mb", }, expectedOk: true, }, { name: "map with negative float64 value", input: map[string]interface{}{ "value": float64(-123.456), "unit": "mb", }, expectedValue: &common.MeteringValue{ Value: 0, // negative values are not converted, leaving Value as 0 Unit: "mb", }, expectedOk: true, }, { name: "map with string value (unsupported)", input: map[string]interface{}{ "value": "100", "unit": "ru", }, expectedValue: nil, expectedOk: false, }, { name: "map missing value field", input: map[string]interface{}{ "unit": "ru", }, expectedValue: nil, expectedOk: false, }, { name: "map missing unit field", input: map[string]interface{}{ "value": uint64(100), }, expectedValue: nil, expectedOk: false, }, { name: "map with non-string unit", input: map[string]interface{}{ "value": uint64(100), "unit": 123, }, expectedValue: nil, expectedOk: false, }, { name: "non-map interface", input: "invalid", expectedValue: nil, expectedOk: false, }, { name: "slice input (unsupported)", input: []string{"a", "b"}, expectedValue: nil, expectedOk: false, }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { result, ok := generator.ConvertToMeteringValue(tt.input) assert.Equal(t, tt.expectedOk, ok, "Expected ok=%v, got %v", tt.expectedOk, ok) if tt.expectedValue == nil { assert.Nil(t, result) } else { require.NotNil(t, result) assert.Equal(t, tt.expectedValue.Value, result.Value) assert.Equal(t, tt.expectedValue.Unit, result.Unit) } }) } } // TestGenerateResourceManagerSKUs_ComprehensiveScenarios tests various scenarios for resource manager SKUs func TestGenerateResourceManagerSKUs_ComprehensiveScenarios(t *testing.T) { generator := newTestSKUGenerator(t) now := time.Now() enrichmentData := map[string]interface{}{ "minRcu": "1000", "maxRcu": "5000", } t.Run("only OLAP RU data", func(t *testing.T) { dataPoint := map[string]interface{}{ models.NextGenMeteringSourceName: "tikv", models.NextGenMeteringOLAPRU: map[string]interface{}{ "value": uint64(2000), "unit": "ru", }, // No OLTP RU data } records, err := generator.GenerateResourceManagerSKUs( dataPoint, "cluster-001", "pool-001", 5000, 1700000, "aws", "us-east-1", "premium", enrichmentData, nil, now, ) require.NoError(t, err) assert.Len(t, records, 1) // One combined record assert.Contains(t, records[0].Unit, "all:total") // Changed to "all" type assert.Contains(t, records[0].Labels, "\"intervalValue\":\"2000\"") assert.Contains(t, records[0].Labels, "\"olapRU\":\"2000\"") // OLAP value in labels }) t.Run("only OLTP RU data", func(t *testing.T) { dataPoint := map[string]interface{}{ models.NextGenMeteringSourceName: "pd", models.NextGenMeteringOLTPRU: map[string]interface{}{ "value": uint64(1500), "unit": "RU", }, // No OLAP RU data } records, err := generator.GenerateResourceManagerSKUs( dataPoint, "cluster-002", "", 6000, 1700001, "gcp", "us-central1", "basic", enrichmentData, nil, now, ) require.NoError(t, err) assert.Len(t, records, 1) // One combined record assert.Contains(t, records[0].Unit, "all:total") // Changed to "all" type assert.Contains(t, records[0].Labels, "\"intervalValue\":\"1500\"") assert.Contains(t, records[0].Labels, "\"oltpRU\":\"1500\"") // OLTP value in labels }) t.Run("both OLAP and OLTP RU with cross-AZ traffic", func(t *testing.T) { dataPoint := map[string]interface{}{ models.NextGenMeteringSourceName: "tikv", models.NextGenMeteringOLAPRU: map[string]interface{}{ "value": uint64(3000), "unit": "ru", }, models.NextGenMeteringOLTPRU: map[string]interface{}{ "value": uint64(2500), "unit": "RU", }, models.NextGenMeteringCrossAZTraffic: map[string]interface{}{ "value": uint64(1048576), // 1MB "unit": "bytes", }, } records, err := generator.GenerateResourceManagerSKUs( dataPoint, "cluster-003", "pool-003", 7000, 1700002, "azure", "eastus", "premium", enrichmentData, nil, now, ) require.NoError(t, err) assert.Len(t, records, 2) // Combined RCU + CrossAZ (now only 2 records) // Verify combined RCU record var rcuRecord *models.BillingReportV2Deltalake for i := range records { if strings.Contains(records[i].Unit, "all:total") { rcuRecord = &records[i] break } } require.NotNil(t, rcuRecord) assert.Contains(t, rcuRecord.Labels, "\"intervalValue\":\"5500\"") // Total: 3000 + 2500 = 5500 assert.Contains(t, rcuRecord.Labels, "\"olapRU\":\"3000\"") // OLAP in labels assert.Contains(t, rcuRecord.Labels, "\"oltpRU\":\"2500\"") // OLTP in labels // Verify CrossAZ record var crossAzRecord *models.BillingReportV2Deltalake for i := range records { if strings.Contains(records[i].Unit, "data-transfer-cross-az") { crossAzRecord = &records[i] break } } require.NotNil(t, crossAzRecord) assert.Equal(t, "1048576", crossAzRecord.UnitValue) }) t.Run("zero values should still report RCU", func(t *testing.T) { dataPoint := map[string]interface{}{ models.NextGenMeteringSourceName: "tikv", models.NextGenMeteringOLAPRU: map[string]interface{}{ "value": uint64(0), "unit": "ru", }, models.NextGenMeteringOLTPRU: map[string]interface{}{ "value": uint64(0), "unit": "RU", }, models.NextGenMeteringCrossAZTraffic: map[string]interface{}{ "value": uint64(0), "unit": "bytes", }, } records, err := generator.GenerateResourceManagerSKUs( dataPoint, "cluster-004", "", 8000, 1700003, "aws", "us-west-1", "premium", enrichmentData, nil, now, ) require.NoError(t, err) // Changed: totalRU is always reported even if 0, but crossAZ traffic with 0 is skipped assert.Len(t, records, 1) // One RCU record (totalRU=0 is reported) assert.Contains(t, records[0].Unit, "all:total") assert.Contains(t, records[0].Labels, "\"intervalValue\":\"0\"") }) t.Run("invalid metering values should still report with defaults", func(t *testing.T) { dataPoint := map[string]interface{}{ models.NextGenMeteringSourceName: "tikv", models.NextGenMeteringOLAPRU: "invalid-data", // Invalid type models.NextGenMeteringOLTPRU: nil, // Nil value models.NextGenMeteringCrossAZTraffic: map[string]interface{}{ // Missing required fields "unit": "bytes", }, } records, err := generator.GenerateResourceManagerSKUs( dataPoint, "cluster-005", "", 9000, 1700004, "gcp", "europe-west1", "premium", enrichmentData, nil, now, ) require.NoError(t, err) // Changed: Invalid values are treated as 0, but totalRU is still reported assert.Len(t, records, 1) // One RCU record with totalRU=0 assert.Contains(t, records[0].Unit, "all:total") assert.Contains(t, records[0].Labels, "\"intervalValue\":\"0\"") }) t.Run("default source name when missing", func(t *testing.T) { dataPoint := map[string]interface{}{ // Missing source_name - should default to "pd" models.NextGenMeteringOLAPRU: map[string]interface{}{ "value": uint64(1000), "unit": "ru", }, } records, err := generator.GenerateResourceManagerSKUs( dataPoint, "cluster-006", "", 10000, 1700005, "aws", "eu-west-1", "premium", enrichmentData, nil, now, ) require.NoError(t, err) assert.Len(t, records, 1) // The source name should be defaulted to "pd" internally assert.NotEmpty(t, records[0].Unit) }) } // TestGenerateProxySKUs_ComprehensiveScenarios tests various scenarios for proxy SKUs func TestGenerateProxySKUs_ComprehensiveScenarios(t *testing.T) { generator := newTestSKUGenerator(t) now := time.Now() enrichmentData := map[string]interface{}{} t.Run("only outbound bytes", func(t *testing.T) { dataPoint := map[string]interface{}{ models.NextGenMeteringSourceName: "tiproxy", models.NextGenMeteringOutBoundBytes: map[string]interface{}{ "value": uint64(5242880), // 5MB "unit": "bytes", }, } records, err := generator.GenerateProxySKUs( dataPoint, "cluster-proxy-001", "proxy-pool-001", 11000, 1700010, "aws", "us-west-2", "premium", enrichmentData, now, ) require.NoError(t, err) assert.Len(t, records, 1) assert.Contains(t, records[0].Unit, "data-transfer-out") assert.Equal(t, "5242880", records[0].UnitValue) assert.Contains(t, records[0].Labels, "\"sourceUnit\":\"bytes\"") }) t.Run("only private outbound bytes", func(t *testing.T) { dataPoint := map[string]interface{}{ models.NextGenMeteringSourceName: "tiproxy", models.NextGenMeteringPrivateOutBoundBytes: map[string]interface{}{ "value": uint64(3145728), // 3MB "unit": "bytes", }, } records, err := generator.GenerateProxySKUs( dataPoint, "cluster-proxy-001", "proxy-pool-001", 11000, 1700010, "aws", "us-west-2", "premium", enrichmentData, now, ) require.NoError(t, err) assert.Len(t, records, 1) assert.Equal(t, "data-transfer-out:aws:us-west-2:premium:tiproxy:private", records[0].Unit) assert.Equal(t, "3145728", records[0].UnitValue) assert.Contains(t, records[0].Labels, "\"sourceUnit\":\"bytes\"") }) t.Run("only public outbound bytes", func(t *testing.T) { dataPoint := map[string]interface{}{ models.NextGenMeteringSourceName: "tiproxy", models.NextGenMeteringPublicOutBoundBytes: map[string]interface{}{ "value": uint64(7340032), // 7MB "unit": "bytes", }, } records, err := generator.GenerateProxySKUs( dataPoint, "cluster-proxy-001", "proxy-pool-001", 11000, 1700010, "aws", "us-west-2", "premium", enrichmentData, now, ) require.NoError(t, err) assert.Len(t, records, 1) assert.Equal(t, "data-transfer-out:aws:us-west-2:premium:tiproxy:public", records[0].Unit) assert.Equal(t, "7340032", records[0].UnitValue) assert.Contains(t, records[0].Labels, "\"sourceUnit\":\"bytes\"") }) t.Run("only cross zone bytes", func(t *testing.T) { dataPoint := map[string]interface{}{ models.NextGenMeteringSourceName: "proxy", models.NextGenMeteringCrossZoneBytes: map[string]interface{}{ "value": uint64(2097152), // 2MB "unit": "Bytes", }, } records, err := generator.GenerateProxySKUs( dataPoint, "cluster-proxy-002", "", 12000, 1700011, "gcp", "asia-east1", "basic", enrichmentData, now, ) require.NoError(t, err) assert.Len(t, records, 1) assert.Contains(t, records[0].Unit, "data-transfer-cross-az") assert.Equal(t, "2097152", records[0].UnitValue) assert.Contains(t, records[0].Labels, "\"sourceUnit\":\"bytes\"") }) t.Run("old version: outbound with cross zone bytes", func(t *testing.T) { dataPoint := map[string]interface{}{ models.NextGenMeteringSourceName: "gateway", models.NextGenMeteringOutBoundBytes: map[string]interface{}{ "value": uint64(10485760), // 10MB "unit": "BYTES", }, models.NextGenMeteringCrossZoneBytes: map[string]interface{}{ "value": uint64(3145728), // 3MB "unit": "bytes", }, } records, err := generator.GenerateProxySKUs( dataPoint, "cluster-proxy-003", "proxy-pool-003", 13000, 1700012, "azure", "westeurope", "premium", enrichmentData, now, ) require.NoError(t, err) assert.Len(t, records, 2) // old version outbound + cross-zone // Verify outbound record var outboundRecord *models.BillingReportV2Deltalake for i := range records { if strings.Contains(records[i].Unit, "data-transfer-out") { outboundRecord = &records[i] break } } require.NotNil(t, outboundRecord) assert.Equal(t, "10485760", outboundRecord.UnitValue) // Verify cross-zone record var crossZoneRecord *models.BillingReportV2Deltalake for i := range records { if strings.Contains(records[i].Unit, "data-transfer-cross-az") { crossZoneRecord = &records[i] break } } require.NotNil(t, crossZoneRecord) assert.Equal(t, "3145728", crossZoneRecord.UnitValue) }) t.Run("new version: both private and public outbound with cross zone bytes", func(t *testing.T) { dataPoint := map[string]interface{}{ models.NextGenMeteringSourceName: "gateway", models.NextGenMeteringPrivateOutBoundBytes: map[string]interface{}{ "value": uint64(4194304), // 4MB "unit": "bytes", }, models.NextGenMeteringPublicOutBoundBytes: map[string]interface{}{ "value": uint64(6291456), // 6MB "unit": "bytes", }, models.NextGenMeteringCrossZoneBytes: map[string]interface{}{ "value": uint64(2097152), // 2MB "unit": "bytes", }, } records, err := generator.GenerateProxySKUs( dataPoint, "cluster-proxy-004", "proxy-pool-004", 14000, 1700013, "aws", "us-east-1", "premium", enrichmentData, now, ) require.NoError(t, err) assert.Len(t, records, 3) // private outbound + public outbound + cross-zone // Build unit map for easier verification unitMap := make(map[string]string) for _, record := range records { unitMap[record.Unit] = record.UnitValue } assert.Equal(t, "4194304", unitMap["data-transfer-out:aws:us-east-1:premium:gateway:private"]) assert.Equal(t, "6291456", unitMap["data-transfer-out:aws:us-east-1:premium:gateway:public"]) assert.Equal(t, "2097152", unitMap["data-transfer-cross-az:aws:us-east-1:premium:gateway"]) }) t.Run("zero values should be skipped", func(t *testing.T) { dataPoint := map[string]interface{}{ models.NextGenMeteringSourceName: "proxy", models.NextGenMeteringOutBoundBytes: map[string]interface{}{ "value": uint64(0), "unit": "bytes", }, models.NextGenMeteringCrossZoneBytes: map[string]interface{}{ "value": uint64(0), "unit": "bytes", }, } records, err := generator.GenerateProxySKUs( dataPoint, "cluster-proxy-004", "", 14000, 1700013, "aws", "ap-southeast-1", "premium", enrichmentData, now, ) require.NoError(t, err) assert.Len(t, records, 0) // No records for zero values }) t.Run("invalid metering values should be handled", func(t *testing.T) { dataPoint := map[string]interface{}{ models.NextGenMeteringSourceName: "proxy", models.NextGenMeteringOutBoundBytes: "invalid-data", models.NextGenMeteringCrossZoneBytes: nil, } records, err := generator.GenerateProxySKUs( dataPoint, "cluster-proxy-005", "", 15000, 1700014, "gcp", "us-east4", "premium", enrichmentData, now, ) require.NoError(t, err) assert.Len(t, records, 0) // No valid records }) t.Run("default source name when missing", func(t *testing.T) { dataPoint := map[string]interface{}{ // Missing source_name - should default to "proxy" models.NextGenMeteringOutBoundBytes: map[string]interface{}{ "value": uint64(1048576), // 1MB "unit": "bytes", }, } records, err := generator.GenerateProxySKUs( dataPoint, "cluster-proxy-006", "", 16000, 1700015, "azure", "northeurope", "basic", enrichmentData, now, ) require.NoError(t, err) assert.Len(t, records, 1) // Should use default source name internally assert.NotEmpty(t, records[0].Unit) }) } func TestGenerateRCUV2SKU(t *testing.T) { generator := newTestSKUGenerator(t) now := time.Now() clusterID := "test-cluster-123" sharedPoolID := "pool-456" tenantID := uint64(789) timestamp := now.Unix() cloudProvider := "aws" region := "us-west-2" plan := "premium" component := "pd" totalRUV2 := uint64(15000) ruUnit := "RU" tidbRUV2 := uint64(5000) tikvRUV2 := uint64(6000) tiflashRUV2 := uint64(4000) enrichmentData := map[string]interface{}{ models.NextGenMetaMaxRCU: "20000", } record, err := generator.GenerateRCUV2SKU(clusterID, sharedPoolID, tenantID, timestamp, cloudProvider, region, plan, component, totalRUV2, ruUnit, tidbRUV2, tikvRUV2, tiflashRUV2, enrichmentData, now) require.NoError(t, err) assert.Equal(t, "v2", record.SchemeVersion) assert.Equal(t, tenantID, record.TenantID) assert.NotEmpty(t, record.CostID) assert.Equal(t, timestamp, record.CostTimestamp) assert.Contains(t, record.Unit, "rcu:aws:us-west-2:premium:all:total") assert.Equal(t, "1", record.UnitValue) assert.Equal(t, int64(1), record.Cost) assert.Equal(t, models.NextGenSKUTypePublic, record.SKUType) var labels map[string]string err = json.Unmarshal([]byte(record.Labels), &labels) require.NoError(t, err) assert.Equal(t, "ru", labels[models.NextGenSKULabelKeySourceUnit]) assert.Equal(t, clusterID, labels[models.NextGenSKULabelKeyClusterID]) assert.Equal(t, sharedPoolID, labels[models.NextGenSKULabelKeySharedPoolID]) assert.Equal(t, "5000", labels[models.NextGenSKULabelKeyTiDBRUV2]) assert.Equal(t, "6000", labels[models.NextGenSKULabelKeyTiKVRUV2]) assert.Equal(t, "4000", labels[models.NextGenSKULabelKeyTiFlashRUV2]) assert.Equal(t, "15000", labels[models.NextGenSKULabelKeyIntervalValue]) assert.Equal(t, models.NextgenFixMinRU, labels[models.NextGenSKULabelKeyIntervalMin]) assert.Equal(t, "20000", labels[models.NextGenSKULabelKeyIntervalMax]) } func TestGenerateRCUV2SKU_SystemCluster(t *testing.T) { generator := newTestSKUGenerator(t) now := time.Now() record, err := generator.GenerateRCUV2SKU( models.NextGenMeteringSystemID, "", 789, now.Unix(), "aws", "us-west-2", "premium", "pd", 15000, "RU", 5000, 6000, 4000, map[string]interface{}{}, now, ) require.NoError(t, err) assert.Equal(t, models.NextGenSKUTypePrivate, record.SKUType) } func TestGenerateRCUV2SKU_EmptyCluster(t *testing.T) { generator := newTestSKUGenerator(t) now := time.Now() record, err := generator.GenerateRCUV2SKU( "", "", 789, now.Unix(), "aws", "us-west-2", "premium", "pd", 15000, "RU", 5000, 6000, 4000, map[string]interface{}{}, now, ) require.NoError(t, err) assert.Equal(t, models.NextGenSKUTypePrivate, record.SKUType) } func TestGenerateArchiveRCUSKU(t *testing.T) { generator := newTestSKUGenerator(t) now := time.Now() clusterID := "test-cluster-123" record, err := generator.GenerateArchiveRCUSKU( clusterID, "", 789, now.Unix(), "aws", "us-west-2", "premium", "pd", 15000, "RU", 10000, 5000, map[string]interface{}{}, now, ) require.NoError(t, err) assert.Equal(t, models.NextGenSKUTypePrivate, record.SKUType, "archive SKU should always be private") assert.Contains(t, record.Unit, "rcu:aws:us-west-2:premium:all:total") } func TestGenerateArchiveRCUSKU_SystemCluster(t *testing.T) { generator := newTestSKUGenerator(t) now := time.Now() record, err := generator.GenerateArchiveRCUSKU( models.NextGenMeteringSystemID, "", 789, now.Unix(), "aws", "us-west-2", "premium", "pd", 15000, "RU", 10000, 5000, map[string]interface{}{}, now, ) require.NoError(t, err) assert.Equal(t, models.NextGenSKUTypePrivate, record.SKUType) } func TestGenerateVCPUUsageSKU(t *testing.T) { generator := newTestSKUGenerator(t) now := time.Now() t.Run("basic vCPU SKU generation", func(t *testing.T) { vcpuValue := uint64(16) dataPoint := map[string]interface{}{ "usage": map[string]interface{}{ "value": vcpuValue, "unit": "vCPU", }, } records, err := generator.GenerateVCPUUsageSKUs( dataPoint, "", // clusterID - empty for BYOC "shared-pool-123", // sharedPoolID 6000, // tenantID 1700000, // timestamp "aws", // cloudProvider "us-west-2", // region "byoc", // plan map[string]interface{}{}, now, ) require.NoError(t, err) require.Len(t, records, 1) record := records[0] assert.Equal(t, "vcpu:aws:us-west-2:byoc:total", record.Unit) assert.Equal(t, uint64(6000), record.TenantID) assert.Equal(t, int64(1700000), record.CostTimestamp) assert.Equal(t, "16", record.UnitValue) assert.Equal(t, int64(1), record.Cost) assert.Equal(t, models.NextGenSKUTypePublic, record.SKUType) assert.Equal(t, "v2", record.SchemeVersion) assert.NotEmpty(t, record.CostID) var labels map[string]string require.NoError(t, json.Unmarshal([]byte(record.Labels), &labels)) assert.Equal(t, "vcpu", labels[models.NextGenSKULabelKeySourceUnit]) assert.Equal(t, "shared-pool-123", labels[models.NextGenSKULabelKeySharedPoolID]) assert.Equal(t, "16", labels[models.NextGenSKULabelKeyIntervalValue]) _, hasClusterID := labels[models.NextGenSKULabelKeyClusterID] assert.False(t, hasClusterID, "clusterID should not be present when empty") }) t.Run("zero vCPU value still generates record", func(t *testing.T) { dataPoint := map[string]interface{}{ "usage": map[string]interface{}{ "value": uint64(0), "unit": "vCPU", }, } records, err := generator.GenerateVCPUUsageSKUs( dataPoint, "", "pool-456", 7000, 1700001, "gcp", "us-central1", "byoc", map[string]interface{}{}, now, ) require.NoError(t, err) require.Len(t, records, 1) assert.Equal(t, "0", records[0].UnitValue) }) t.Run("missing vCPU field returns empty records", func(t *testing.T) { dataPoint := map[string]interface{}{} records, err := generator.GenerateVCPUUsageSKUs( dataPoint, "", "pool-789", 8000, 1700002, "aws", "us-east-1", "byoc", map[string]interface{}{}, now, ) require.NoError(t, err) assert.Empty(t, records) }) t.Run("invalid vCPU field returns empty records", func(t *testing.T) { dataPoint := map[string]interface{}{ "usage": "not-a-number", } records, err := generator.GenerateVCPUUsageSKUs( dataPoint, "", "pool-000", 9000, 1700003, "aws", "us-east-1", "byoc", map[string]interface{}{}, now, ) require.NoError(t, err) assert.Empty(t, records) }) } func TestGenerateResourceManagerSKUs_RCUV2(t *testing.T) { generator := newTestSKUGenerator(t) now := time.Now() timestamp := now.Unix() t.Run("ru_version=2 with ru_v2 present", func(t *testing.T) { dataPoint := map[string]interface{}{ models.NextGenMeteringOLAPRU: map[string]interface{}{ "value": uint64(10000), "unit": "RU", }, models.NextGenMeteringOLTPRU: map[string]interface{}{ "value": uint64(5000), "unit": "RU", }, models.NextGenMeteringTiDBRUV2: map[string]interface{}{ "value": uint64(3000), "unit": "RU", }, models.NextGenMeteringTiKVRUV2: map[string]interface{}{ "value": uint64(3000), "unit": "RU", }, models.NextGenMeteringTiFlashRUV2: map[string]interface{}{ "value": uint64(2000), "unit": "RU", }, } enrichmentData := map[string]interface{}{ models.NextGenMetaRUVersion: "2", models.NextGenMetaMaxRCU: "50000", } records, err := generator.GenerateResourceManagerSKUs( dataPoint, "cluster-v2", "pool-v2", 100, timestamp, "aws", "us-west-2", "premium", enrichmentData, nil, now, ) require.NoError(t, err) assert.Len(t, records, 2, "v2 scenario should generate exactly 2 records: rcuv2 + archive rcu") // Record 0: rcu with V2 labels (public) assert.Contains(t, records[0].Unit, "rcu:") assert.NotContains(t, records[0].Unit, "rcuv2:") assert.Equal(t, models.NextGenSKUTypePublic, records[0].SKUType) var v2Labels map[string]string require.NoError(t, json.Unmarshal([]byte(records[0].Labels), &v2Labels)) assert.Equal(t, "8000", v2Labels[models.NextGenSKULabelKeyIntervalValue]) // 3000 + 3000 + 2000 assert.Equal(t, "3000", v2Labels[models.NextGenSKULabelKeyTiDBRUV2]) assert.Equal(t, "3000", v2Labels[models.NextGenSKULabelKeyTiKVRUV2]) assert.Equal(t, "2000", v2Labels[models.NextGenSKULabelKeyTiFlashRUV2]) // Record 1: archive rcu (private) assert.Contains(t, records[1].Unit, "rcu:") assert.NotContains(t, records[1].Unit, "rcuv2:") assert.Equal(t, models.NextGenSKUTypePrivate, records[1].SKUType) var archiveLabels map[string]string require.NoError(t, json.Unmarshal([]byte(records[1].Labels), &archiveLabels)) assert.Equal(t, "15000", archiveLabels[models.NextGenSKULabelKeyIntervalValue]) // 10000 + 5000 }) t.Run("ru_version=2 with all ru_v2 missing should fallback to regular rcu", func(t *testing.T) { dataPoint := map[string]interface{}{ models.NextGenMeteringOLAPRU: map[string]interface{}{ "value": uint64(10000), "unit": "RU", }, models.NextGenMeteringOLTPRU: map[string]interface{}{ "value": uint64(5000), "unit": "RU", }, // No tidb_ru_v2, tikv_ru_v2, tiflash_ru_v2 } enrichmentData := map[string]interface{}{ models.NextGenMetaRUVersion: "2", } records, err := generator.GenerateResourceManagerSKUs( dataPoint, "cluster-v2-fallback", "", 100, timestamp, "aws", "us-west-2", "premium", enrichmentData, nil, now, ) require.NoError(t, err) assert.Len(t, records, 1, "v2 with all ru_v2 missing should fallback to 1 regular rcu") assert.Contains(t, records[0].Unit, "rcu:") assert.Equal(t, models.NextGenSKUTypePublic, records[0].SKUType) }) t.Run("ru_version=2 with partial ru_v2 should still generate rcuv2", func(t *testing.T) { dataPoint := map[string]interface{}{ models.NextGenMeteringOLAPRU: map[string]interface{}{ "value": uint64(10000), "unit": "RU", }, models.NextGenMeteringOLTPRU: map[string]interface{}{ "value": uint64(5000), "unit": "RU", }, models.NextGenMeteringTiDBRUV2: map[string]interface{}{ "value": uint64(7000), "unit": "RU", }, // tikv_ru_v2 and tiflash_ru_v2 missing } enrichmentData := map[string]interface{}{ models.NextGenMetaRUVersion: "2", models.NextGenMetaMaxRCU: "50000", } records, err := generator.GenerateResourceManagerSKUs( dataPoint, "cluster-v2-partial", "", 100, timestamp, "aws", "us-west-2", "premium", enrichmentData, nil, now, ) require.NoError(t, err) assert.Len(t, records, 2, "v2 with partial ru_v2 should generate 2 records") assert.Contains(t, records[0].Unit, "rcu:") assert.NotContains(t, records[0].Unit, "rcuv2:") var v2Labels map[string]string require.NoError(t, json.Unmarshal([]byte(records[0].Labels), &v2Labels)) assert.Equal(t, "7000", v2Labels[models.NextGenSKULabelKeyIntervalValue]) // only tidb_ru_v2 assert.Equal(t, "7000", v2Labels[models.NextGenSKULabelKeyTiDBRUV2]) assert.Equal(t, "0", v2Labels[models.NextGenSKULabelKeyTiKVRUV2]) assert.Equal(t, "0", v2Labels[models.NextGenSKULabelKeyTiFlashRUV2]) }) t.Run("no ru_version should generate regular rcu only", func(t *testing.T) { dataPoint := map[string]interface{}{ models.NextGenMeteringOLAPRU: map[string]interface{}{ "value": uint64(10000), "unit": "RU", }, models.NextGenMeteringOLTPRU: map[string]interface{}{ "value": uint64(5000), "unit": "RU", }, } enrichmentData := map[string]interface{}{} records, err := generator.GenerateResourceManagerSKUs( dataPoint, "cluster-old", "", 100, timestamp, "aws", "us-west-2", "premium", enrichmentData, nil, now, ) require.NoError(t, err) assert.Len(t, records, 1) assert.Contains(t, records[0].Unit, "rcu:") assert.NotContains(t, records[0].Unit, "rcuv2:") assert.Equal(t, models.NextGenSKUTypePublic, records[0].SKUType) }) t.Run("ru_version=1 should generate regular rcu only", func(t *testing.T) { dataPoint := map[string]interface{}{ models.NextGenMeteringOLAPRU: map[string]interface{}{ "value": uint64(10000), "unit": "RU", }, models.NextGenMeteringOLTPRU: map[string]interface{}{ "value": uint64(5000), "unit": "RU", }, } enrichmentData := map[string]interface{}{ models.NextGenMetaRUVersion: "1", } records, err := generator.GenerateResourceManagerSKUs( dataPoint, "cluster-v1", "", 100, timestamp, "aws", "us-west-2", "premium", enrichmentData, nil, now, ) require.NoError(t, err) assert.Len(t, records, 1) assert.Contains(t, records[0].Unit, "rcu:") assert.NotContains(t, records[0].Unit, "rcuv2:") }) t.Run("ru_version=2 SYSTEM cluster should generate private rcu with V2 labels", func(t *testing.T) { dataPoint := map[string]interface{}{ models.NextGenMeteringOLAPRU: map[string]interface{}{ "value": uint64(10000), "unit": "RU", }, models.NextGenMeteringOLTPRU: map[string]interface{}{ "value": uint64(5000), "unit": "RU", }, models.NextGenMeteringTiDBRUV2: map[string]interface{}{ "value": uint64(3000), "unit": "RU", }, models.NextGenMeteringTiKVRUV2: map[string]interface{}{ "value": uint64(3000), "unit": "RU", }, models.NextGenMeteringTiFlashRUV2: map[string]interface{}{ "value": uint64(2000), "unit": "RU", }, } enrichmentData := map[string]interface{}{ models.NextGenMetaRUVersion: "2", } records, err := generator.GenerateResourceManagerSKUs( dataPoint, models.NextGenMeteringSystemID, "", 100, timestamp, "aws", "us-west-2", "premium", enrichmentData, nil, now, ) require.NoError(t, err) assert.Len(t, records, 2) assert.Equal(t, models.NextGenSKUTypePrivate, records[0].SKUType, "SYSTEM rcuv2 should be private") assert.Equal(t, models.NextGenSKUTypePrivate, records[1].SKUType, "SYSTEM archive rcu should be private") }) } func TestGenerateResourceManagerSKUs_TopologyDrivenSkipsRCU(t *testing.T) { generator := newTestSKUGenerator(t) now := time.Now() dataPoint := map[string]interface{}{ models.NextGenMeteringOLAPRU: map[string]interface{}{ "value": uint64(1000), "unit": "RU", }, models.NextGenMeteringOLTPRU: map[string]interface{}{ "value": uint64(2000), "unit": "RU", }, models.NextGenMeteringCrossAZTraffic: map[string]interface{}{ "value": uint64(4096), "unit": "bytes", }, models.NextGenMeteringIncrementalWriteByte: map[string]interface{}{ "value": uint64(1024), "unit": "bytes", }, } enrichmentData := map[string]interface{}{ models.NextGenMetaCapacityMode: "TOPOLOGY_DRIVEN", } backupMetaData := map[string]interface{}{ models.NextGenMetaBackupStorageV2: "T2", } records, err := generator.GenerateResourceManagerSKUs( dataPoint, "cluster-topology", "pool-topology", 100, now.Unix(), "aws", "us-west-2", "reserved", enrichmentData, backupMetaData, now, ) require.NoError(t, err) require.Len(t, records, 2) for _, record := range records { assert.NotContains(t, record.Unit, "rcu:") } assert.Equal(t, "data-transfer-cross-az:aws:us-west-2:reserved:pd", records[0].Unit) assert.Equal(t, "backup-storage:aws:us-west-2:reserved:T2", records[1].Unit) } func TestGenerateReservedNodeSKUs(t *testing.T) { generator := newTestSKUGenerator(t) now := time.Now() dataPoint := map[string]interface{}{ models.NextGenMeteringNodeType: "tikv", "usage": map[string]interface{}{ "value": uint64(16), "unit": "vCPU", }, } records, err := generator.GenerateReservedNodeSKUs( dataPoint, "cluster-1", "pool-1", 100, now.Unix(), "aws", "us-west-2", "premium_reserved", nil, now, ) require.NoError(t, err) require.Len(t, records, 1) record := records[0] assert.Equal(t, "vcpu:aws:us-west-2:reserved:all", record.Unit) assert.Equal(t, "16", record.UnitValue) assert.Equal(t, models.NextGenSKUTypePublic, record.SKUType) var labels map[string]string require.NoError(t, json.Unmarshal([]byte(record.Labels), &labels)) assert.Equal(t, "vcpu", labels[models.NextGenSKULabelKeySourceUnit]) assert.Equal(t, "tikv", labels[models.NextGenSKULabelKeyNodeType]) assert.Equal(t, "cluster-1", labels[models.NextGenSKULabelKeyClusterID]) assert.Equal(t, "pool-1", labels[models.NextGenSKULabelKeySharedPoolID]) }