Addition of tests for dlqMessageHandler in common domain

common/domain/dlqMessageHandler.go → common/domain/dlq_message_handler.go

@@ -18,17 +18,17 @@
 // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 // THE SOFTWARE.
 
-//go:generate mockgen -package $GOPACKAGE -source $GOFILE -destination dlqMessageHandler_mock.go
+//go:generate mockgen -package $GOPACKAGE -source $GOFILE -destination dlq_message_handler_mock.go
 
 package domain
 
 import (
 	"context"
 	"sync"
 	"sync/atomic"
-	"time"
 
 	"github.com/uber/cadence/common"
+	"github.com/uber/cadence/common/clock"
 	"github.com/uber/cadence/common/log"
 	"github.com/uber/cadence/common/log/tag"
 	"github.com/uber/cadence/common/metrics"
@@ -51,6 +51,7 @@ type (
 		replicationQueue   ReplicationQueue
 		logger             log.Logger
 		metricsClient      metrics.Client
+		timeSrc            clock.TimeSource
 		done               chan struct{}
 		status             int32
 
@@ -65,12 +66,14 @@ func NewDLQMessageHandler(
 	replicationQueue ReplicationQueue,
 	logger log.Logger,
 	metricsClient metrics.Client,
+	timeSource clock.TimeSource,
 ) DLQMessageHandler {
 	return &dlqMessageHandlerImpl{
 		replicationHandler: replicationHandler,
 		replicationQueue:   replicationQueue,
 		logger:             logger,
 		metricsClient:      metricsClient,
+		timeSrc:            timeSource,
 		done:               make(chan struct{}),
 		lastCount:          -1,
 	}
@@ -213,14 +216,14 @@ func (d *dlqMessageHandlerImpl) Merge(
 }
 
 func (d *dlqMessageHandlerImpl) emitDLQSizeMetricsLoop() {
-	ticker := time.NewTicker(queueSizeQueryInterval)
+	ticker := d.timeSrc.NewTicker(queueSizeQueryInterval)
 	defer ticker.Stop()
 
 	for {
 		select {
 		case <-d.done:
 			return
-		case <-ticker.C:
+		case <-ticker.Chan():
 			err := d.fetchAndEmitDLQSize(context.Background())
 			if err != nil {
 				d.logger.Warn("Failed to get DLQ size.", tag.Error(err))

common/domain/dlqMessageHandler_test.go → common/domain/dlq_message_handler_test.go

@@ -22,14 +22,21 @@ package domain
 
 import (
 	"context"
+	"errors"
 	"fmt"
+	"sync/atomic"
 	"testing"
+	"time"
 
 	"github.com/golang/mock/gomock"
 	"github.com/pborman/uuid"
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"github.com/stretchr/testify/suite"
 
+	"github.com/uber/cadence/common"
+	"github.com/uber/cadence/common/clock"
+	"github.com/uber/cadence/common/log/tag"
 	"github.com/uber/cadence/common/log/testlogger"
 	"github.com/uber/cadence/common/metrics"
 	"github.com/uber/cadence/common/types"
@@ -73,6 +80,7 @@ func (s *dlqMessageHandlerSuite) SetupTest() {
 		s.mockReplicationQueue,
 		logger,
 		metrics.NewNoopMetricsClient(),
+		clock.NewMockedTimeSource(),
 	).(*dlqMessageHandlerImpl)
 }
 
@@ -102,6 +110,146 @@ func (s *dlqMessageHandlerSuite) TestReadMessages() {
 	s.Nil(token)
 }
 
+func TestDLQMessageHandler_Start(t *testing.T) {
+	tests := []struct {
+		name           string
+		setupMocks     func(m *MockReplicationQueue)
+		initialStatus  int32
+		expectedStatus int32
+	}{
+		{
+			name: "Should start when initialized",
+			setupMocks: func(m *MockReplicationQueue) {
+				m.EXPECT().GetDLQSize(gomock.Any()).Return(int64(1), nil).AnyTimes()
+			},
+			initialStatus:  common.DaemonStatusInitialized,
+			expectedStatus: common.DaemonStatusStarted,
+		},
+		{
+			name: "Should not start when already started",
+			setupMocks: func(m *MockReplicationQueue) {
+				// No calls expected since the handler should recognize it's already started
+			},
+			initialStatus:  common.DaemonStatusStarted,
+			expectedStatus: common.DaemonStatusStarted,
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			ctrl := gomock.NewController(t)
+			defer ctrl.Finish()
+
+			mockQueue := NewMockReplicationQueue(ctrl)
+			mockReplicationTaskExecutor := NewMockReplicationTaskExecutor(ctrl)
+			mockLogger := testlogger.New(t)
+			mockedTime := clock.NewMockedTimeSource()
+			tc.setupMocks(mockQueue)
+
+			doneChan := make(chan struct{})
+			handler := &dlqMessageHandlerImpl{
+				replicationHandler: mockReplicationTaskExecutor,
+				replicationQueue:   mockQueue,
+				metricsClient:      metrics.NewNoopMetricsClient(),
+				logger:             mockLogger,
+				done:               doneChan,
+				status:             tc.initialStatus,
+				timeSrc:            mockedTime,
+			}
+
+			handler.Start()
+			mockedTime.Advance(5 * time.Minute)
+			defer handler.Stop()
+			assert.Equal(t, tc.expectedStatus, atomic.LoadInt32(&handler.status), assert.Equal(t, tc.expectedStatus, atomic.LoadInt32(&handler.status), "Handler did not reach the expected status for test case: %s", tc.name))
+		})
+	}
+}
+
+func (s *dlqMessageHandlerSuite) TestStop() {
+	tests := []struct {
+		name           string
+		initialStatus  int32
+		expectedStatus int32
+		shouldStop     bool
+	}{
+		{
+			name:           "Should stop when started",
+			initialStatus:  common.DaemonStatusStarted,
+			expectedStatus: common.DaemonStatusStopped,
+			shouldStop:     true,
+		},
+		{
+			name:           "Should not stop when not started",
+			initialStatus:  common.DaemonStatusInitialized,
+			expectedStatus: common.DaemonStatusInitialized,
+			shouldStop:     false,
+		},
+	}
+
+	for _, test := range tests {
+		s.Run(test.name, func() {
+			atomic.StoreInt32(&s.dlqMessageHandler.status, test.initialStatus)
+			s.dlqMessageHandler.Stop()
+			if test.shouldStop {
+				s.dlqMessageHandler.logger.Info("Domain DLQ handler shutting down.")
+			}
+			s.Equal(test.expectedStatus, atomic.LoadInt32(&s.dlqMessageHandler.status))
+		})
+	}
+}
+
+func (s *dlqMessageHandlerSuite) TestCount() {
+	tests := []struct {
+		name          string
+		forceFetch    bool
+		lastCount     int64
+		fetchSize     int64
+		fetchError    error
+		expectedCount int64
+		expectedError error
+	}{
+		{
+			name:          "Force fetch with error",
+			forceFetch:    true,
+			lastCount:     10,
+			fetchSize:     0,
+			fetchError:    fmt.Errorf("fetch error"),
+			expectedCount: 0,
+			expectedError: fmt.Errorf("fetch error"),
+		},
+		{
+			name:          "Force fetch with success",
+			forceFetch:    true,
+			lastCount:     10,
+			fetchSize:     20,
+			fetchError:    nil,
+			expectedCount: 20,
+			expectedError: nil,
+		},
+		{
+			name:          "No fetch needed",
+			forceFetch:    false,
+			lastCount:     30,
+			expectedCount: 30,
+			expectedError: nil,
+		},
+	}
+
+	for _, test := range tests {
+		s.Run(test.name, func() {
+			s.mockReplicationQueue.EXPECT().GetDLQSize(gomock.Any()).Return(test.fetchSize, test.fetchError).MaxTimes(1)
+			s.dlqMessageHandler.lastCount = test.lastCount
+			count, err := s.dlqMessageHandler.Count(context.Background(), test.forceFetch)
+			s.Equal(test.expectedCount, count)
+			if test.expectedError != nil {
+				s.Equal(test.expectedError, err)
+			} else {
+				s.NoError(err)
+			}
+		})
+	}
+}
+
 func (s *dlqMessageHandlerSuite) TestReadMessages_ThrowErrorOnGetDLQAckLevel() {
 	lastMessageID := int64(20)
 	pageSize := 100
@@ -359,3 +507,49 @@ func (s *dlqMessageHandlerSuite) TestMergeMessages_IgnoreErrorOnUpdateDLQAckLeve
 	s.NoError(err)
 	s.Nil(token)
 }
+
+func (s *dlqMessageHandlerSuite) TestMergeMessages_NonDomainTask() {
+	ackLevel := int64(10)
+	lastMessageID := int64(20)
+	pageSize := 100
+	pageToken := []byte{}
+
+	// Create a task that mimics a non-domain replication task by not setting DomainTaskAttributes
+	tasks := []*types.ReplicationTask{
+		{
+			TaskType:     types.ReplicationTaskTypeDomain.Ptr(), // Still set to domain but no attributes
+			SourceTaskID: 1,
+		},
+	}
+
+	s.mockReplicationQueue.EXPECT().GetDLQAckLevel(gomock.Any()).Return(ackLevel, nil).Times(1)
+	s.mockReplicationQueue.EXPECT().GetMessagesFromDLQ(gomock.Any(), ackLevel, lastMessageID, pageSize, pageToken).
+		Return(tasks, nil, nil).Times(1)
+
+	token, err := s.dlqMessageHandler.Merge(context.Background(), lastMessageID, pageSize, pageToken)
+
+	s.NotNil(err)
+	s.IsType(&types.InternalServiceError{}, err)
+	s.Equal("Encounter non domain replication task in domain replication queue.", err.Error())
+	s.Nil(token)
+}
+
+func (s *dlqMessageHandlerSuite) TestEmitDLQSizeMetricsLoop_ErrorHandling() {
+	expectedError := fmt.Errorf("error fetching DLQ size")
+	s.mockReplicationQueue.EXPECT().GetDLQSize(gomock.Any()).Return(int64(0), expectedError).AnyTimes()
+
+	// Start the metrics loop in a goroutine
+	go s.dlqMessageHandler.emitDLQSizeMetricsLoop()
+
+	// Allow some time for the goroutine to run and tick at least once
+	time.Sleep(1 * time.Nanosecond)
+
+	// Close the done channel to signal the loop to stop
+	close(s.dlqMessageHandler.done)
+
+	// Wait a bit to ensure the loop exits
+	time.Sleep(1 * time.Nanosecond)
+
+	s.dlqMessageHandler.logger.Warn("Failed to get DLQ size.", tag.Error(errors.New("DomainReplicationQueueSizeLimit")))
+
+}

service/frontend/admin/handler.go

@@ -36,6 +36,7 @@ import (
 	"github.com/uber/cadence/common/asyncworkflow/queueconfigapi"
 	"github.com/uber/cadence/common/backoff"
 	"github.com/uber/cadence/common/client"
+	"github.com/uber/cadence/common/clock"
 	"github.com/uber/cadence/common/codec"
 	"github.com/uber/cadence/common/definition"
 	"github.com/uber/cadence/common/domain"
@@ -131,6 +132,7 @@ func NewHandler(
 			resource.GetDomainReplicationQueue(),
 			resource.GetLogger(),
 			resource.GetMetricsClient(),
+			clock.NewRealTimeSource(),
 		),
 		domainFailoverWatcher: domain.NewFailoverWatcher(
 			resource.GetDomainCache(),