package handler

import (
	"context"
	"errors"
	"testing"

	"github.com/stretchr/testify/assert"
	mock_facade "github.com/tidbcloud/observability/o11ymetering/mock/pkg/facade"
	"github.com/tidbcloud/observability/o11ymetering/pkg/common/constants"
	"github.com/tidbcloud/observability/o11ymetering/pkg/common/models"
	"github.com/tidbcloud/observability/o11ymetering/pkg/facade"
	"github.com/tidbcloud/observability/pkg/common/configs"
	o11ystructs "github.com/tidbcloud/observability/pkg/common/structs"
	"go.uber.org/mock/gomock"
)

// setupMocks initializes gomock controller and mocks for facade interfaces.
func setupMocks(t *testing.T) (*gomock.Controller, *mock_facade.MockO11yMetaQuery, *mock_facade.MockO11yMetricsQuery) {
	ctrl := gomock.NewController(t)
	mockMetaQuery := mock_facade.NewMockO11yMetaQuery(ctrl)
	mockMetricsQuery := mock_facade.NewMockO11yMetricsQuery(ctrl)
	facade.SetO11yMetaQuery(mockMetaQuery)
	facade.SetO11yMetricsQuery(mockMetricsQuery)
	return ctrl, mockMetaQuery, mockMetricsQuery
}

// DedicatedRegionalControlMetricCollection Tests
func TestDedicatedRegionalControlMetricCollection(t *testing.T) {
	configs.Init(o11ystructs.ServiceInfo{
		Region: "us-east-1",
		Vendor: "aws",
	})
	t.Run("Build success", func(t *testing.T) {
		h := &DedicatedRegionalControlMetricCollection{Env: constants.Prod}
		task, err := h.Build(context.Background(), "2023-01-01 00:00:00")
		assert.NoError(t, err)
		assert.Equal(t, constants.MetricsCollection, task.StatisticClass)
	})

	t.Run("Execute success", func(t *testing.T) {
		ctrl, mockMetaQuery, mockMetricsQuery := setupMocks(t)
		defer ctrl.Finish()

		h := &DedicatedRegionalControlMetricCollection{Env: constants.Prod}
		task := models.StatisticTask{ID: "1", Env: constants.Prod, Period: "2023-01-01 00:00:00"}
		mockCluster := o11ystructs.ClusterBaseInfo{ID: "seed-cluster-id"}
		mockMetricsResp := &models.O11yMetricsResponse{Data: models.O11yMetrics{Result: []models.O11yMetricsResult{
			{O11yMetricsLabelSet: models.O11yMetricsLabelSet{Tenant: "t1"}, Value: []interface{}{float64(1), "100"}},
		}}}

		mockMetaQuery.EXPECT().GetDedicatedSeedCluster(gomock.Any(), constants.Prod).Return(mockCluster, nil)
		mockMetricsQuery.EXPECT().QueryO11yMetricsInstant(gomock.Any(), gomock.Any(), gomock.Any(), mockCluster.ID, gomock.Any(), gomock.Any(), constants.Prod).Return(mockMetricsResp, nil)

		records, err := h.Execute(context.Background(), task)
		assert.NoError(t, err)
		assert.Len(t, records, 1)
		assert.Equal(t, 100.0, records[0].Value)
	})

	t.Run("Execute GetDedicatedSeedCluster error", func(t *testing.T) {
		ctrl, mockMetaQuery, _ := setupMocks(t)
		defer ctrl.Finish()

		h := &DedicatedRegionalControlMetricCollection{Env: constants.Prod}
		task := models.StatisticTask{Env: constants.Prod}

		mockMetaQuery.EXPECT().GetDedicatedSeedCluster(gomock.Any(), constants.Prod).Return(o11ystructs.ClusterBaseInfo{}, errors.New("cluster error"))

		_, err := h.Execute(context.Background(), task)
		assert.Error(t, err)
		assert.Contains(t, err.Error(), "cluster error")
	})

	t.Run("Execute QueryO11yMetricsInstant error", func(t *testing.T) {
		ctrl, mockMetaQuery, mockMetricsQuery := setupMocks(t)
		defer ctrl.Finish()

		h := &DedicatedRegionalControlMetricCollection{Env: constants.Prod}
		task := models.StatisticTask{Env: constants.Prod, Period: "2023-01-01 00:00:00"}
		mockCluster := o11ystructs.ClusterBaseInfo{ID: "seed-cluster-id"}

		mockMetaQuery.EXPECT().GetDedicatedSeedCluster(gomock.Any(), constants.Prod).Return(mockCluster, nil)
		mockMetricsQuery.EXPECT().QueryO11yMetricsInstant(gomock.Any(), gomock.Any(), gomock.Any(), mockCluster.ID, gomock.Any(), gomock.Any(), constants.Prod).Return(nil, errors.New("query error"))

		_, err := h.Execute(context.Background(), task)
		assert.Error(t, err)
		assert.Contains(t, err.Error(), "query error")
	})
}

// DedicatedDataplaneS3Collection Tests
func TestDedicatedDataplaneS3Collection(t *testing.T) {
	configs.Init(o11ystructs.ServiceInfo{
		Region: "us-east-1",
		Vendor: "aws",
	})
	t.Run("Build success", func(t *testing.T) {
		h := &DedicatedDataplaneS3Collection{Env: constants.Prod}
		task, err := h.Build(context.Background(), "2023-01-01 00:00:00")
		assert.NoError(t, err)
		assert.Equal(t, constants.LogsStorageS3, task.StatisticClass)
	})

	t.Run("Execute success", func(t *testing.T) {
		ctrl, mockMetaQuery, mockMetricsQuery := setupMocks(t)
		defer ctrl.Finish()

		h := &DedicatedDataplaneS3Collection{Env: constants.Prod}
		task := models.StatisticTask{ID: "1", Env: constants.Prod, Period: "2023-01-01 00:00:00"}
		mockInfraClusters := []o11ystructs.ClusterBaseInfo{
			{ID: "infra-cluster-0"},
		}
		mockClusters := []o11ystructs.ClusterBaseInfo{{ID: "c1"}, {ID: "c2"}}
		mockMetricsResp := &models.O11yMetricsResponse{Data: models.O11yMetrics{Result: []models.O11yMetricsResult{
			{O11yMetricsLabelSet: models.O11yMetricsLabelSet{Tenant: "t1"}, Value: []interface{}{float64(1), "100"}},
		}}}

		mockMetaQuery.EXPECT().ListDedicatedK8sInfraClusters(gomock.Any(), constants.Prod).Return(mockInfraClusters, nil)
		for _, queryTemplate := range dedicatedDataplaneS3LessLabelQueryMap {
			query := renderS3ComponentType(queryTemplate)
			mockMetricsQuery.EXPECT().QueryO11yMetricsInstant(gomock.Any(), gomock.Any(), gomock.Any(), "infra-cluster-0", gomock.Any(), query, constants.Prod).Return(mockMetricsResp, nil)
		}

		mockMetaQuery.EXPECT().ListO11yDedicatedClusters(gomock.Any(), constants.Prod).Return(mockClusters, nil)
		for _, cluster := range mockClusters {
			for _, queryTemplate := range dedicatedDataplaneS3LessLabelQueryMap {
				renderS3ComponentType(queryTemplate)
				mockMetricsQuery.EXPECT().QueryO11yMetricsInstant(gomock.Any(), gomock.Any(), gomock.Any(), cluster.ID, gomock.Any(), gomock.Any(), constants.Prod).Return(mockMetricsResp, nil)
			}
		}

		records, err := h.Execute(context.Background(), task)
		assert.NoError(t, err)
		assert.Len(t, records, (len(mockClusters)+len(mockInfraClusters))*len(dedicatedDataplaneS3LessLabelQueryMap))
	})

}

// DedicatedRegionalControlS3Collection Tests
func TestDedicatedRegionalControlS3Collection(t *testing.T) {
	configs.Init(o11ystructs.ServiceInfo{
		Region: "us-east-1",
		Vendor: "aws",
	})
	t.Run("Build success", func(t *testing.T) {
		h := &DedicatedRegionalControlS3Collection{Env: constants.Prod}
		task, err := h.Build(context.Background(), "2023-01-01 00:00:00")
		assert.NoError(t, err)
		assert.Equal(t, constants.LogsStorageS3, task.StatisticClass)
	})

	t.Run("Execute success", func(t *testing.T) {
		ctrl, mockMetaQuery, mockMetricsQuery := setupMocks(t)
		defer ctrl.Finish()

		h := &DedicatedRegionalControlS3Collection{Env: constants.Prod}
		task := models.StatisticTask{ID: "1", Env: constants.Prod, Period: "2023-01-01 00:00:00"}
		mockCluster := o11ystructs.ClusterBaseInfo{ID: "seed-cluster-id"}
		mockMetricsResp := &models.O11yMetricsResponse{Data: models.O11yMetrics{Result: []models.O11yMetricsResult{
			{O11yMetricsLabelSet: models.O11yMetricsLabelSet{Tenant: "t1"}, Value: []interface{}{float64(1), "100"}},
		}}}

		mockMetaQuery.EXPECT().GetDedicatedSeedCluster(gomock.Any(), constants.Prod).Return(mockCluster, nil)
		for _, queryTemplate := range dedicatedRegionalControlS3QueryMap {
			query := renderS3ComponentType(queryTemplate)
			mockMetricsQuery.EXPECT().QueryO11yMetricsInstant(gomock.Any(), gomock.Any(), gomock.Any(), mockCluster.ID, gomock.Any(), query, constants.Prod).Return(mockMetricsResp, nil)
		}

		records, err := h.Execute(context.Background(), task)
		assert.NoError(t, err)
		assert.Len(t, records, len(dedicatedRegionalControlS3QueryMap))
	})
}

// DedicatedRegionalControlLokiCollection Tests
func TestDedicatedRegionalControlLokiCollection(t *testing.T) {
	configs.Init(o11ystructs.ServiceInfo{
		Region: "us-east-1",
		Vendor: "aws",
	})
	t.Run("Build success", func(t *testing.T) {
		h := &DedicatedRegionalControlLokiCollection{Env: constants.Prod}
		task, err := h.Build(context.Background(), "2023-01-01 00:00:00")
		assert.NoError(t, err)
		assert.Equal(t, constants.LogsCollectionLoki, task.StatisticClass)
	})

	t.Run("Execute success", func(t *testing.T) {
		ctrl, mockMetaQuery, mockMetricsQuery := setupMocks(t)
		defer ctrl.Finish()

		h := &DedicatedRegionalControlLokiCollection{Env: constants.Prod}
		task := models.StatisticTask{ID: "1", Env: constants.Prod, Period: "2023-01-01 00:00:00"}
		mockCluster := o11ystructs.ClusterBaseInfo{ID: "seed-cluster-id"}
		mockMetricsResp := &models.O11yMetricsResponse{Data: models.O11yMetrics{Result: []models.O11yMetricsResult{
			{O11yMetricsLabelSet: models.O11yMetricsLabelSet{Tenant: "t1"}, Value: []interface{}{float64(1), "100"}},
		}}}

		mockMetaQuery.EXPECT().GetDedicatedSeedCluster(gomock.Any(), constants.Prod).Return(mockCluster, nil)
		for _, queryTemplate := range dedicatedRegionalControlQueryMap {
			mockMetricsQuery.EXPECT().QueryO11yMetricsInstant(gomock.Any(), gomock.Any(), gomock.Any(), mockCluster.ID, gomock.Any(), queryTemplate, constants.Prod).Return(mockMetricsResp, nil)
		}

		records, err := h.Execute(context.Background(), task)
		assert.NoError(t, err)
		assert.Len(t, records, len(dedicatedRegionalControlQueryMap))
	})
}

// ServerlessMetricCollection Tests
func TestServerlessMetricCollection(t *testing.T) {
	configs.Init(o11ystructs.ServiceInfo{
		Region: "us-east-1",
		Vendor: "aws",
	})
	t.Run("Build success", func(t *testing.T) {
		h := &ServerlessMetricCollection{Env: constants.Prod}
		task, err := h.Build(context.Background(), "2023-01-01 00:00:00")
		assert.NoError(t, err)
		assert.Equal(t, constants.MetricsCollection, task.StatisticClass)
	})

	t.Run("Execute success", func(t *testing.T) {
		ctrl, mockMetaQuery, mockMetricsQuery := setupMocks(t)
		defer ctrl.Finish()

		h := &ServerlessMetricCollection{Env: constants.Prod}
		task := models.StatisticTask{ID: "1", Env: constants.Prod, Period: "2023-01-01 00:00:00"}
		mockCluster := o11ystructs.ClusterBaseInfo{ID: "serverless-infra-cluster-id"}
		mockMetricsResp := &models.O11yMetricsResponse{Data: models.O11yMetrics{Result: []models.O11yMetricsResult{
			{O11yMetricsLabelSet: models.O11yMetricsLabelSet{Tenant: "t1"}, Value: []interface{}{float64(1), "100"}},
		}}}

		mockMetaQuery.EXPECT().GetServerlessInfraCluster(gomock.Any(), constants.Prod).Return(mockCluster, nil)
		mockMetricsQuery.EXPECT().QueryO11yMetricsInstant(gomock.Any(), gomock.Any(), gomock.Any(), mockCluster.ID, gomock.Any(), gomock.Any(), constants.Prod).Return(mockMetricsResp, nil)

		records, err := h.Execute(context.Background(), task)
		assert.NoError(t, err)
		assert.Len(t, records, 1)
	})
}

// ServerlessS3Collection Tests
func TestServerlessS3Collection(t *testing.T) {
	configs.Init(o11ystructs.ServiceInfo{
		Region: "us-east-1",
		Vendor: "aws",
	})
	t.Run("Build success", func(t *testing.T) {
		h := &ServerlessS3Collection{Env: constants.Prod}
		task, err := h.Build(context.Background(), "2023-01-01 00:00:00")
		assert.NoError(t, err)
		assert.Equal(t, constants.LogsStorageS3, task.StatisticClass)
	})

	t.Run("Execute success", func(t *testing.T) {
		ctrl, mockMetaQuery, mockMetricsQuery := setupMocks(t)
		defer ctrl.Finish()

		h := &ServerlessS3Collection{Env: constants.Prod}
		task := models.StatisticTask{ID: "1", Env: constants.Prod, Period: "2023-01-01 00:00:00"}
		mockCluster := o11ystructs.ClusterBaseInfo{ID: "serverless-infra-cluster-id"}
		mockMetricsResp := &models.O11yMetricsResponse{Data: models.O11yMetrics{Result: []models.O11yMetricsResult{
			{O11yMetricsLabelSet: models.O11yMetricsLabelSet{Tenant: "t1"}, Value: []interface{}{float64(1), "100"}},
		}}}

		mockMetaQuery.EXPECT().GetServerlessInfraCluster(gomock.Any(), constants.Prod).Return(mockCluster, nil)
		for _, queryTemplate := range serverlessS3QueryMap {
			query := renderS3ComponentType(queryTemplate)
			mockMetricsQuery.EXPECT().QueryO11yMetricsInstant(gomock.Any(), gomock.Any(), gomock.Any(), mockCluster.ID, gomock.Any(), query, constants.Prod).Return(mockMetricsResp, nil)
		}

		records, err := h.Execute(context.Background(), task)
		assert.NoError(t, err)
		assert.Len(t, records, len(serverlessS3QueryMap))
	})
}

// ServerlessLokiCollection Tests
func TestServerlessLokiCollection(t *testing.T) {
	configs.Init(o11ystructs.ServiceInfo{
		Region: "us-east-1",
		Vendor: "aws",
	})
	t.Run("Build success", func(t *testing.T) {
		h := &ServerlessLokiCollection{Env: constants.Prod}
		task, err := h.Build(context.Background(), "2023-01-01 00:00:00")
		assert.NoError(t, err)
		assert.Equal(t, constants.LogsCollectionLoki, task.StatisticClass)
	})

	t.Run("Execute success", func(t *testing.T) {
		ctrl, mockMetaQuery, mockMetricsQuery := setupMocks(t)
		defer ctrl.Finish()

		h := &ServerlessLokiCollection{Env: constants.Prod}
		task := models.StatisticTask{ID: "1", Env: constants.Prod, Period: "2023-01-01 00:00:00"}
		mockCluster := o11ystructs.ClusterBaseInfo{ID: "serverless-infra-cluster-id"}
		mockMetricsResp := &models.O11yMetricsResponse{Data: models.O11yMetrics{Result: []models.O11yMetricsResult{
			{O11yMetricsLabelSet: models.O11yMetricsLabelSet{Tenant: "t1"}, Value: []interface{}{float64(1), "100"}},
		}}}

		mockMetaQuery.EXPECT().GetServerlessInfraCluster(gomock.Any(), constants.Prod).Return(mockCluster, nil)
		for _, queryTemplate := range serverlessQueryMap {
			mockMetricsQuery.EXPECT().QueryO11yMetricsInstant(gomock.Any(), gomock.Any(), gomock.Any(), mockCluster.ID, gomock.Any(), queryTemplate, constants.Prod).Return(mockMetricsResp, nil)
		}

		records, err := h.Execute(context.Background(), task)
		assert.NoError(t, err)
		assert.Len(t, records, len(serverlessQueryMap))
	})
}

// DedicatedDataplaneMetricCollection Tests
func TestDedicatedDataplaneMetricCollection(t *testing.T) {
	configs.Init(o11ystructs.ServiceInfo{
		Region: "us-east-1",
		Vendor: "aws",
	})
	t.Run("Build success", func(t *testing.T) {
		h := &DedicatedDataplaneMetricCollection{Env: constants.Prod}     // Initialize handler
		task, err := h.Build(context.Background(), "2023-01-01 00:00:00") // Call Build method
		assert.NoError(t, err)                                            // Assert no error
		assert.Equal(t, constants.MetricsCollection, task.StatisticClass) // Verify StatisticClass
		assert.Equal(t, constants.DedicatedDataplane, task.MeteringGroup) // Verify MeteringGroup
	})

	t.Run("Execute success", func(t *testing.T) {
		ctrl, mockMetaQuery, mockMetricsQuery := setupMocks(t) // Setup mocks
		defer ctrl.Finish()                                    // Ensure mocks are satisfied

		h := &DedicatedDataplaneMetricCollection{Env: constants.Prod}                             // Initialize handler
		task := models.StatisticTask{ID: "1", Env: constants.Prod, Period: "2023-01-01 00:00:00"} // Create a sample task
		mockInfraClusters := []o11ystructs.ClusterBaseInfo{
			{ID: "infra-cluster-0"},
			{ID: "infra-cluster-1"},
		}
		mockDataplaneCluster := o11ystructs.ClusterBaseInfo{ID: "dataplane-cluster-id"} // Mock dataplane cluster for topsql

		mockMetricsResp := &models.O11yMetricsResponse{ // Mock metrics response
			Data: models.O11yMetrics{
				Result: []models.O11yMetricsResult{
					{O11yMetricsLabelSet: models.O11yMetricsLabelSet{Tenant: "t1"}, Value: []interface{}{float64(1), "100"}},
				},
			},
		}

		mockMetaQuery.EXPECT().ListDedicatedK8sInfraClusters(gomock.Any(), constants.Prod).Return(mockInfraClusters, nil)                                                                               // Expect GetDedicatedSeedCluster call
		mockMetricsQuery.EXPECT().QueryO11yMetricsInstant(gomock.Any(), gomock.Any(), gomock.Any(), "infra-cluster-0", gomock.Any(), dedicatedMetricQuery, constants.Prod).Return(mockMetricsResp, nil) // Expect QueryO11yMetricsInstant for common metrics
		mockMetricsQuery.EXPECT().QueryO11yMetricsInstant(gomock.Any(), gomock.Any(), gomock.Any(), "infra-cluster-1", gomock.Any(), dedicatedMetricQuery, constants.Prod).Return(mockMetricsResp, nil) // Expect QueryO11yMetricsInstant for common metrics

		mockMetaQuery.EXPECT().ListO11yDedicatedClusters(gomock.Any(), constants.Prod).Return([]o11ystructs.ClusterBaseInfo{mockDataplaneCluster}, nil)                                                             // Expect ListO11yDedicatedClusters for topsql
		mockMetricsQuery.EXPECT().QueryO11yMetricsInstant(gomock.Any(), gomock.Any(), gomock.Any(), mockDataplaneCluster.ID, gomock.Any(), dedicatedMetricTopsqlQuery, constants.Prod).Return(mockMetricsResp, nil) // Expect QueryO11yMetricsInstant for topsql

		records, err := h.Execute(context.Background(), task) // Execute the method under test
		assert.NoError(t, err)                                // Assert no error
		assert.Len(t, records, 3)                             // Expect 3 records (2 from common, 1 from topsql)
		assert.Equal(t, 100.0, records[0].Value)              // Assert value
		assert.Equal(t, 100.0, records[1].Value)              // Assert value
		assert.Equal(t, 100.0, records[2].Value)              // Assert value
	})

	t.Run("Execute ListDedicatedK8sInfraClusters error", func(t *testing.T) {
		ctrl, mockMetaQuery, _ := setupMocks(t) // Setup mocks
		defer ctrl.Finish()                     // Ensure mocks are satisfied

		h := &DedicatedDataplaneMetricCollection{Env: constants.Prod} // Initialize handler
		task := models.StatisticTask{Env: constants.Prod}             // Create a sample task

		mockMetaQuery.EXPECT().ListDedicatedK8sInfraClusters(gomock.Any(), constants.Prod).Return(nil, errors.New("infra clusters error")) // Mock GetDedicatedSeedCluster to return an error

		_, err := h.Execute(context.Background(), task)         // Execute the method under test
		assert.Error(t, err)                                    // Assert error
		assert.Contains(t, err.Error(), "infra clusters error") // Assert error message
	})

	t.Run("Execute common metrics query error", func(t *testing.T) {
		ctrl, mockMetaQuery, mockMetricsQuery := setupMocks(t) // Setup mocks
		defer ctrl.Finish()                                    // Ensure mocks are satisfied

		h := &DedicatedDataplaneMetricCollection{Env: constants.Prod}                    // Initialize handler
		task := models.StatisticTask{Env: constants.Prod, Period: "2023-01-01 00:00:00"} // Create a sample task
		mockInfraClusters := []o11ystructs.ClusterBaseInfo{
			{ID: "infra-cluster-0"},
		}

		mockMetaQuery.EXPECT().ListDedicatedK8sInfraClusters(gomock.Any(), constants.Prod).Return(mockInfraClusters, nil)                                                                                                        // Expect GetDedicatedSeedCluster call
		mockMetricsQuery.EXPECT().QueryO11yMetricsInstant(gomock.Any(), gomock.Any(), gomock.Any(), "infra-cluster-0", gomock.Any(), dedicatedMetricQuery, constants.Prod).Return(nil, errors.New("common metrics query error")) // Mock common metrics query to return error

		_, err := h.Execute(context.Background(), task)               // Execute the method under test
		assert.Error(t, err)                                          // Assert error
		assert.Contains(t, err.Error(), "common metrics query error") // Assert error message
	})

	t.Run("AppendPerClusterTopsql ListO11yDedicatedClusters error", func(t *testing.T) {
		ctrl, mockMetaQuery, mockMetricsQuery := setupMocks(t) // Setup mocks
		defer ctrl.Finish()                                    // Ensure mocks are satisfied

		h := &DedicatedDataplaneMetricCollection{Env: constants.Prod}                    // Initialize handler
		task := models.StatisticTask{Env: constants.Prod, Period: "2023-01-01 00:00:00"} // Create a sample task
		mockInfraClusters := []o11ystructs.ClusterBaseInfo{
			{ID: "infra-cluster-0"},
		}
		mockMetricsResp := &models.O11yMetricsResponse{Data: models.O11yMetrics{Result: []models.O11yMetricsResult{
			{O11yMetricsLabelSet: models.O11yMetricsLabelSet{Tenant: "t1"}, Value: []interface{}{float64(1), "100"}},
		}}}

		mockMetaQuery.EXPECT().ListDedicatedK8sInfraClusters(gomock.Any(), constants.Prod).Return(mockInfraClusters, nil)                                                                               // Expect GetDedicatedSeedCluster call
		mockMetricsQuery.EXPECT().QueryO11yMetricsInstant(gomock.Any(), gomock.Any(), gomock.Any(), "infra-cluster-0", gomock.Any(), dedicatedMetricQuery, constants.Prod).Return(mockMetricsResp, nil) // Expect QueryO11yMetricsInstant for common metrics

		mockMetaQuery.EXPECT().ListO11yDedicatedClusters(gomock.Any(), constants.Prod).Return(nil, errors.New("list clusters error")) // Mock ListO11yDedicatedClusters to return error

		_, err := h.Execute(context.Background(), task)        // Execute the method under test
		assert.Error(t, err)                                   // Assert error
		assert.Contains(t, err.Error(), "list clusters error") // Assert error message
	})

	t.Run("AppendPerClusterTopsql QueryO11yMetricsInstant error", func(t *testing.T) {
		ctrl, mockMetaQuery, mockMetricsQuery := setupMocks(t) // Setup mocks
		defer ctrl.Finish()                                    // Ensure mocks are satisfied

		h := &DedicatedDataplaneMetricCollection{Env: constants.Prod}                    // Initialize handler
		task := models.StatisticTask{Env: constants.Prod, Period: "2023-01-01 00:00:00"} // Create a sample task
		mockInfraClusters := []o11ystructs.ClusterBaseInfo{
			{ID: "infra-cluster-0"},
		}
		mockDataplaneCluster := o11ystructs.ClusterBaseInfo{ID: "dataplane-cluster-id"} // Mock dataplane cluster for topsql
		mockMetricsResp := &models.O11yMetricsResponse{Data: models.O11yMetrics{Result: []models.O11yMetricsResult{
			{O11yMetricsLabelSet: models.O11yMetricsLabelSet{Tenant: "t1"}, Value: []interface{}{float64(1), "100"}},
		}}}

		mockMetaQuery.EXPECT().ListDedicatedK8sInfraClusters(gomock.Any(), constants.Prod).Return(mockInfraClusters, nil)                                                                               // Expect GetDedicatedSeedCluster call
		mockMetricsQuery.EXPECT().QueryO11yMetricsInstant(gomock.Any(), gomock.Any(), gomock.Any(), "infra-cluster-0", gomock.Any(), dedicatedMetricQuery, constants.Prod).Return(mockMetricsResp, nil) // Expect QueryO11yMetricsInstant for common metrics

		mockMetaQuery.EXPECT().ListO11yDedicatedClusters(gomock.Any(), constants.Prod).Return([]o11ystructs.ClusterBaseInfo{mockDataplaneCluster}, nil)                                                                              // Expect ListO11yDedicatedClusters for topsql
		mockMetricsQuery.EXPECT().QueryO11yMetricsInstant(gomock.Any(), gomock.Any(), gomock.Any(), mockDataplaneCluster.ID, gomock.Any(), dedicatedMetricTopsqlQuery, constants.Prod).Return(nil, errors.New("topsql query error")) // Mock topsql query to return error

		_, err := h.Execute(context.Background(), task)       // Execute the method under test
		assert.Error(t, err)                                  // Assert error
		assert.Contains(t, err.Error(), "topsql query error") // Assert error message
	})
}

// DedicatedDataplaneLokiCollection Tests
func TestDedicatedDataplaneLokiCollection(t *testing.T) {
	configs.Init(o11ystructs.ServiceInfo{
		Region: "us-east-1",
		Vendor: "aws",
	})
	t.Run("Build success", func(t *testing.T) {
		h := &DedicatedDataplaneLokiCollection{Env: constants.Prod}        // Initialize handler
		task, err := h.Build(context.Background(), "2023-01-01 00:00:00")  // Call Build method
		assert.NoError(t, err)                                             // Assert no error
		assert.Equal(t, constants.LogsCollectionLoki, task.StatisticClass) // Verify StatisticClass
		assert.Equal(t, constants.DedicatedDataplane, task.MeteringGroup)  // Verify MeteringGroup
	})

	t.Run("Execute success", func(t *testing.T) {
		ctrl, mockMetaQuery, mockMetricsQuery := setupMocks(t) // Setup mocks
		defer ctrl.Finish()                                    // Ensure mocks are satisfied

		h := &DedicatedDataplaneLokiCollection{Env: constants.Prod}                               // Initialize handler
		task := models.StatisticTask{ID: "1", Env: constants.Prod, Period: "2023-01-01 00:00:00"} // Create a sample task
		mockInfraClusters := []o11ystructs.ClusterBaseInfo{
			{ID: "infra-cluster-0"},
		}
		mockMetricsResp := &models.O11yMetricsResponse{ // Mock metrics response
			Data: models.O11yMetrics{
				Result: []models.O11yMetricsResult{
					{O11yMetricsLabelSet: models.O11yMetricsLabelSet{Tenant: "t1"}, Value: []interface{}{float64(1), "100"}},
				},
			},
		}

		mockMetaQuery.EXPECT().ListDedicatedK8sInfraClusters(gomock.Any(), constants.Prod).Return(mockInfraClusters, nil) // Expect call to GetDedicatedSeedCluster
		for _, queryTemplate := range dedicatedDataplaneLessLabelQueryMap {                                               // Iterate through queries
			// query is queryTemplate, no fmt.Sprintf in this Execute method for now
			mockMetricsQuery.EXPECT().QueryO11yMetricsInstant(
				gomock.Any(), gomock.Any(), gomock.Any(), "infra-cluster-0", gomock.Any(), queryTemplate, constants.Prod,
			).Return(mockMetricsResp, nil) // Expect QueryO11yMetricsInstant for each query
		}

		records, err := h.Execute(context.Background(), task)            // Execute the method under test
		assert.NoError(t, err)                                           // Assert no error
		assert.Len(t, records, len(dedicatedDataplaneLessLabelQueryMap)) // Assert correct number of records
		assert.Equal(t, 100.0, records[0].Value)                         // Assert record value
		assert.Equal(t, constants.Logs, records[0].Category)
	})

	t.Run("Execute ListDedicatedK8sInfraClusters error", func(t *testing.T) {
		ctrl, mockMetaQuery, _ := setupMocks(t) // Setup mocks
		defer ctrl.Finish()                     // Ensure mocks are satisfied

		h := &DedicatedDataplaneLokiCollection{Env: constants.Prod} // Initialize handler
		task := models.StatisticTask{Env: constants.Prod}           // Create a sample task

		mockMetaQuery.EXPECT().ListDedicatedK8sInfraClusters(gomock.Any(), constants.Prod).Return(nil, errors.New("cluster error")) // Mock GetDedicatedSeedCluster to return an error

		_, err := h.Execute(context.Background(), task)  // Execute the method under test
		assert.Error(t, err)                             // Assert error
		assert.Contains(t, err.Error(), "cluster error") // Assert error message
	})

	t.Run("Execute QueryO11yMetricsInstant error", func(t *testing.T) {
		ctrl, mockMetaQuery, mockMetricsQuery := setupMocks(t) // Setup mocks
		defer ctrl.Finish()                                    // Ensure mocks are satisfied

		h := &DedicatedDataplaneLokiCollection{Env: constants.Prod}                      // Initialize handler
		task := models.StatisticTask{Env: constants.Prod, Period: "2023-01-01 00:00:00"} // Create a sample task
		mockInfraClusters := []o11ystructs.ClusterBaseInfo{
			{ID: "infra-cluster-0"},
		}

		mockMetaQuery.EXPECT().ListDedicatedK8sInfraClusters(gomock.Any(), constants.Prod).Return(mockInfraClusters, nil) // Expect call to GetDedicatedSeedCluster
		mockMetricsQuery.EXPECT().QueryO11yMetricsInstant(
			gomock.Any(), gomock.Any(), gomock.Any(), "infra-cluster-0", gomock.Any(), gomock.Any(), constants.Prod,
		).Return(nil, errors.New("query error")) // Mock QueryO11yMetricsInstant to return an error

		_, err := h.Execute(context.Background(), task) // Execute the method under test
		assert.Error(t, err)                            // Assert error
		assert.Contains(t, err.Error(), "query error")  // Assert error message
	})
}

// CloudPlatformBaseS3Collection Tests
func TestCloudPlatformBaseS3Collection(t *testing.T) {
	configs.Init(o11ystructs.ServiceInfo{
		Region: "us-east-1",
		Vendor: "aws",
	})
	t.Run("Build success", func(t *testing.T) {
		h := &CloudPlatformBaseS3Collection{Env: constants.Prod}          // Initialize handler
		task, err := h.Build(context.Background(), "2023-01-01 00:00:00") // Call Build method
		assert.NoError(t, err)                                            // Assert no error
		assert.Equal(t, constants.LogsStorageS3, task.StatisticClass)     // Verify StatisticClass
		assert.Equal(t, constants.CloudPlatform, task.MeteringGroup)      // Verify MeteringGroup
	})

	t.Run("Execute success", func(t *testing.T) {
		ctrl, mockMetaQuery, mockMetricsQuery := setupMocks(t) // Setup mocks
		defer ctrl.Finish()                                    // Ensure mocks are satisfied

		h := &CloudPlatformBaseS3Collection{Env: constants.Prod}                                  // Initialize handler
		task := models.StatisticTask{ID: "1", Env: constants.Prod, Period: "2023-01-01 00:00:00"} // Create a sample task
		mockCluster := o11ystructs.ClusterBaseInfo{ID: "global-base-cluster-id"}                  // Mock cluster info
		mockMetricsResp := &models.O11yMetricsResponse{                                           // Mock metrics response
			Data: models.O11yMetrics{
				Result: []models.O11yMetricsResult{
					{O11yMetricsLabelSet: models.O11yMetricsLabelSet{Tenant: "t1"}, Value: []interface{}{float64(1), "100"}},
				},
			},
		}

		mockMetaQuery.EXPECT().GetGlobalBaseCluster(gomock.Any(), constants.Prod).Return(mockCluster, nil) // Expect call to GetGlobalBaseCluster
		for _, queryTemplate := range baseS3QueryMap {                                                     // Iterate through queries
			query := renderS3ComponentType(queryTemplate) // Render the S3 component type
			mockMetricsQuery.EXPECT().QueryO11yMetricsInstant(
				gomock.Any(), gomock.Any(), gomock.Any(), mockCluster.ID, gomock.Any(), query, constants.Prod,
			).Return(mockMetricsResp, nil) // Expect QueryO11yMetricsInstant for each query
		}

		records, err := h.Execute(context.Background(), task) // Execute the method under test
		assert.NoError(t, err)                                // Assert no error
		assert.Len(t, records, len(baseS3QueryMap))           // Assert correct number of records
		assert.Equal(t, 100.0, records[0].Value)              // Assert record value
		assert.Equal(t, constants.Logs, records[0].Category)
	})

	t.Run("Execute GetGlobalBaseCluster error", func(t *testing.T) {
		ctrl, mockMetaQuery, _ := setupMocks(t) // Setup mocks
		defer ctrl.Finish()                     // Ensure mocks are satisfied

		h := &CloudPlatformBaseS3Collection{Env: constants.Prod} // Initialize handler
		task := models.StatisticTask{Env: constants.Prod}        // Create a sample task

		mockMetaQuery.EXPECT().GetGlobalBaseCluster(gomock.Any(), constants.Prod).Return(o11ystructs.ClusterBaseInfo{}, errors.New("cluster error")) // Mock GetGlobalBaseCluster to return an error

		_, err := h.Execute(context.Background(), task)  // Execute the method under test
		assert.Error(t, err)                             // Assert error
		assert.Contains(t, err.Error(), "cluster error") // Assert error message
	})

	t.Run("Execute QueryO11yMetricsInstant error", func(t *testing.T) {
		ctrl, mockMetaQuery, mockMetricsQuery := setupMocks(t) // Setup mocks
		defer ctrl.Finish()                                    // Ensure mocks are satisfied

		h := &CloudPlatformBaseS3Collection{Env: constants.Prod}                         // Initialize handler
		task := models.StatisticTask{Env: constants.Prod, Period: "2023-01-01 00:00:00"} // Create a sample task
		mockCluster := o11ystructs.ClusterBaseInfo{ID: "global-base-cluster-id"}         // Mock cluster info

		mockMetaQuery.EXPECT().GetGlobalBaseCluster(gomock.Any(), constants.Prod).Return(mockCluster, nil) // Expect call to GetGlobalBaseCluster
		mockMetricsQuery.EXPECT().QueryO11yMetricsInstant(
			gomock.Any(), gomock.Any(), gomock.Any(), mockCluster.ID, gomock.Any(), gomock.Any(), constants.Prod,
		).Return(nil, errors.New("query error")) // Mock QueryO11yMetricsInstant to return an error

		_, err := h.Execute(context.Background(), task) // Execute the method under test
		assert.Error(t, err)                            // Assert error
		assert.Contains(t, err.Error(), "query error")  // Assert error message
	})
}

// CloudPlatformBaseMetricCollection Tests
func TestCloudPlatformBaseMetricCollection(t *testing.T) {
	configs.Init(o11ystructs.ServiceInfo{
		Region: "us-east-1",
		Vendor: "aws",
	})
	t.Run("Build success", func(t *testing.T) {
		h := &CloudPlatformBaseMetricCollection{Env: constants.Prod}      // Initialize handler
		task, err := h.Build(context.Background(), "2023-01-01 00:00:00") // Call Build method
		assert.NoError(t, err)                                            // Assert no error
		assert.Equal(t, constants.MetricsCollection, task.StatisticClass) // Verify StatisticClass
		assert.Equal(t, constants.CloudPlatform, task.MeteringGroup)      // Verify MeteringGroup
	})

	t.Run("Execute success", func(t *testing.T) {
		ctrl, mockMetaQuery, mockMetricsQuery := setupMocks(t) // Setup mocks
		defer ctrl.Finish()                                    // Ensure mocks are satisfied

		h := &CloudPlatformBaseMetricCollection{Env: constants.Prod}                              // Initialize handler
		task := models.StatisticTask{ID: "1", Env: constants.Prod, Period: "2023-01-01 00:00:00"} // Create a sample task
		mockCluster := o11ystructs.ClusterBaseInfo{ID: "global-base-cluster-id"}                  // Mock cluster info
		mockMetricsResp := &models.O11yMetricsResponse{                                           // Mock metrics response
			Data: models.O11yMetrics{
				Result: []models.O11yMetricsResult{
					{O11yMetricsLabelSet: models.O11yMetricsLabelSet{Tenant: "t1"}, Value: []interface{}{float64(1), "100"}},
				},
			},
		}

		mockMetaQuery.EXPECT().GetGlobalBaseCluster(gomock.Any(), constants.Prod).Return(mockCluster, nil) // Expect call to GetGlobalBaseCluster
		mockMetricsQuery.EXPECT().QueryO11yMetricsInstant(
			gomock.Any(), gomock.Any(), gomock.Any(), mockCluster.ID, gomock.Any(), baseQuery, constants.Prod,
		).Return(mockMetricsResp, nil) // Expect QueryO11yMetricsInstant

		records, err := h.Execute(context.Background(), task) // Execute the method under test
		assert.NoError(t, err)                                // Assert no error
		assert.Len(t, records, 1)                             // Assert correct number of records
		assert.Equal(t, 100.0, records[0].Value)              // Assert record value
		assert.Equal(t, constants.Metrics, records[0].Category)
	})

	t.Run("Execute GetGlobalBaseCluster error", func(t *testing.T) {
		ctrl, mockMetaQuery, _ := setupMocks(t) // Setup mocks
		defer ctrl.Finish()                     // Ensure mocks are satisfied

		h := &CloudPlatformBaseMetricCollection{Env: constants.Prod} // Initialize handler
		task := models.StatisticTask{Env: constants.Prod}            // Create a sample task

		mockMetaQuery.EXPECT().GetGlobalBaseCluster(gomock.Any(), constants.Prod).Return(o11ystructs.ClusterBaseInfo{}, errors.New("cluster error")) // Mock GetGlobalBaseCluster to return an error

		_, err := h.Execute(context.Background(), task)  // Execute the method under test
		assert.Error(t, err)                             // Assert error
		assert.Contains(t, err.Error(), "cluster error") // Assert error message
	})

	t.Run("Execute QueryO11yMetricsInstant error", func(t *testing.T) {
		ctrl, mockMetaQuery, mockMetricsQuery := setupMocks(t) // Setup mocks
		defer ctrl.Finish()                                    // Ensure mocks are satisfied

		h := &CloudPlatformBaseMetricCollection{Env: constants.Prod}                     // Initialize handler
		task := models.StatisticTask{Env: constants.Prod, Period: "2023-01-01 00:00:00"} // Create a sample task
		mockCluster := o11ystructs.ClusterBaseInfo{ID: "global-base-cluster-id"}         // Mock cluster info

		mockMetaQuery.EXPECT().GetGlobalBaseCluster(gomock.Any(), constants.Prod).Return(mockCluster, nil) // Expect call to GetGlobalBaseCluster
		mockMetricsQuery.EXPECT().QueryO11yMetricsInstant(
			gomock.Any(), gomock.Any(), gomock.Any(), mockCluster.ID, gomock.Any(), baseQuery, constants.Prod,
		).Return(nil, errors.New("query error")) // Mock QueryO11yMetricsInstant to return an error

		_, err := h.Execute(context.Background(), task) // Execute the method under test
		assert.Error(t, err)                            // Assert error
		assert.Contains(t, err.Error(), "query error")  // Assert error message
	})
}