From c189829d1ea939224a2175f5d1249368cf8f9a71 Mon Sep 17 00:00:00 2001
From: Ilya Ozherelyev <ilya.ozherelyev@uber.com>
Date: Thu, 31 Oct 2024 21:47:00 +0100
Subject: [PATCH] rebase and add tests

---
 internal/internal_task_pollers_test.go | 54 ++++++++++++++++++++++----
 1 file changed, 47 insertions(+), 7 deletions(-)

diff --git a/internal/internal_task_pollers_test.go b/internal/internal_task_pollers_test.go
index 5318ddfd2..91d30efeb 100644
--- a/internal/internal_task_pollers_test.go
+++ b/internal/internal_task_pollers_test.go
@@ -25,6 +25,7 @@ import (
 	"errors"
 	"github.com/golang/mock/gomock"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/require"
 	"github.com/uber-go/tally"
 	"go.uber.org/cadence/.gen/go/cadence/workflowservicetest"
@@ -42,12 +43,6 @@ const (
 	_testIdentity   = "test-worker"
 )
 
-// Enable verbose logging for tests.
-func TestMain(m *testing.M) {
-	EnableVerboseLogging(true)
-	m.Run()
-}
-
 func TestLocalActivityPanic(t *testing.T) {
 	// regression: panics in local activities should not terminate the process
 	s := WorkflowTestSuite{logger: zaptest.NewLogger(t)}
@@ -120,7 +115,6 @@ func TestRespondTaskCompleted_failed(t *testing.T) {
 		assert.NoError(t, err)
 		assert.Nil(t, res)
 	})
-
 }
 
 func TestRespondTaskCompleted_Unsupported(t *testing.T) {
@@ -131,6 +125,52 @@ func TestRespondTaskCompleted_Unsupported(t *testing.T) {
 	})
 }
 
+func TestProcessTask_failures(t *testing.T) {
+	t.Run("shutdown", func(t *testing.T) {
+		poller, _, _, _ := buildWorkflowTaskPoller(t)
+		ch := make(chan struct{})
+		poller.shutdownC = ch
+		close(ch)
+
+		err := poller.ProcessTask(&workflowTask{})
+		assert.ErrorIs(t, err, errShutdown)
+	})
+	t.Run("unsupported task type", func(t *testing.T) {
+		poller, _, _, _ := buildWorkflowTaskPoller(t)
+		assert.Panics(t, func() {
+			_ = poller.ProcessTask(10)
+		})
+	})
+	t.Run("nil task", func(t *testing.T) {
+		poller, _, _, _ := buildWorkflowTaskPoller(t)
+
+		err := poller.ProcessTask(&workflowTask{})
+		assert.NoError(t, err)
+	})
+	t.Run("heartbeat error", func(t *testing.T) {
+		poller, _, mockedTaskHandler, _ := buildWorkflowTaskPoller(t)
+		hearbeatErr := &decisionHeartbeatError{}
+		mockedTaskHandler.EXPECT().ProcessWorkflowTask(mock.Anything, mock.Anything).Return(nil, hearbeatErr)
+		err := poller.ProcessTask(&workflowTask{
+			task: &s.PollForDecisionTaskResponse{},
+		})
+		assert.ErrorIs(t, err, hearbeatErr)
+	})
+	t.Run("ResetStickyTaskList fail", func(t *testing.T) {
+		poller, client, _, _ := buildWorkflowTaskPoller(t)
+		client.EXPECT().ResetStickyTaskList(gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, assert.AnError)
+		err := poller.ProcessTask(&resetStickinessTask{
+			task: &s.ResetStickyTaskListRequest{
+				Execution: &s.WorkflowExecution{
+					WorkflowId: common.StringPtr("test-workflow-id"),
+					RunId:      common.StringPtr("test-run-id"),
+				},
+			},
+		})
+		assert.ErrorIs(t, err, assert.AnError)
+	})
+}
+
 func buildWorkflowTaskPoller(t *testing.T) (*workflowTaskPoller, *workflowservicetest.MockClient, *MockWorkflowTaskHandler, *MockLocalDispatcher) {
 	ctrl := gomock.NewController(t)
 	mockService := workflowservicetest.NewMockClient(ctrl)