Add unit tests for NotifyActivityTask batching and generation

Author: rkannan82

common/testing/testvars/test_vars.go

@@ -396,6 +396,15 @@ func (tv *TestVars) WorkerIdentity() string {
 	return getOrCreate(tv, "worker_identity", tv.uniqueString, tv.stringNSetter)
 }
 
+func (tv *TestVars) WorkerInstanceKey() string {
+	return getOrCreate(tv, "worker_instance_key", tv.uniqueString, tv.stringNSetter)
+}
+
+// ControlQueueName returns the Nexus task queue name used to deliver control tasks to this worker.
+func (tv *TestVars) ControlQueueName(ns string) string {
+	return fmt.Sprintf("/temporal-sys/worker-commands/%s/%s", ns, tv.WorkerInstanceKey())
+}
+
 func (tv *TestVars) TimerID() string {
 	return getOrCreate(tv, "timer_id", tv.uniqueString, tv.stringNSetter)
 }

service/history/api/respondworkflowtaskcompleted/workflow_task_completed_handler_test.go

@@ -384,3 +384,63 @@ func mustMarshalAny(t *testing.T, pb proto.Message) *anypb.Any {
 	require.NoError(t, a.MarshalFrom(pb))
 	return &a
 }
+
+func TestFlushBatchedNotifyActivityTasks(t *testing.T) {
+	t.Parallel()
+
+	t.Run("batches activities by control queue", func(t *testing.T) {
+		ctrl := gomock.NewController(t)
+		ms := historyi.NewMockMutableState(ctrl)
+
+		ms.EXPECT().AddNotifyActivityTasks(
+			[]int64{5, 6, 7},
+			"control-queue-1",
+			gomock.Any(),
+		).Return(nil).Times(1)
+
+		handler := &workflowTaskCompletedHandler{
+			mutableState: ms,
+			pendingActivityCancelsByControlQueue: map[string][]int64{
+				"control-queue-1": {5, 6, 7},
+			},
+		}
+
+		err := handler.flushBatchedNotifyActivityTasks()
+		require.NoError(t, err)
+	})
+
+	t.Run("creates separate tasks for different control queues", func(t *testing.T) {
+		ctrl := gomock.NewController(t)
+		ms := historyi.NewMockMutableState(ctrl)
+
+		ms.EXPECT().AddNotifyActivityTasks(
+			gomock.Any(),
+			gomock.Any(),
+			gomock.Any(),
+		).Return(nil).Times(2)
+
+		handler := &workflowTaskCompletedHandler{
+			mutableState: ms,
+			pendingActivityCancelsByControlQueue: map[string][]int64{
+				"control-queue-1": {5, 6},
+				"control-queue-2": {7, 8},
+			},
+		}
+
+		err := handler.flushBatchedNotifyActivityTasks()
+		require.NoError(t, err)
+	})
+
+	t.Run("does nothing when no pending cancels", func(t *testing.T) {
+		ctrl := gomock.NewController(t)
+		ms := historyi.NewMockMutableState(ctrl)
+
+		handler := &workflowTaskCompletedHandler{
+			mutableState:                         ms,
+			pendingActivityCancelsByControlQueue: nil,
+		}
+
+		err := handler.flushBatchedNotifyActivityTasks()
+		require.NoError(t, err)
+	})
+}

service/history/workflow/task_generator_test.go

@@ -1066,3 +1066,80 @@ func TestTaskGeneratorImpl_GenerateDeleteHistoryEventTask_ActivityRetention(t *t
 		})
 	}
 }
+
+func TestGenerateNotifyActivityTasks(t *testing.T) {
+	t.Parallel()
+
+	testCases := []struct {
+		name              string
+		featureEnabled    bool
+		scheduledEventIDs []int64
+		controlQueue      string
+		expectTask        bool
+	}{
+		{
+			name:              "creates task when enabled with valid inputs",
+			featureEnabled:    true,
+			scheduledEventIDs: []int64{5, 6, 7},
+			controlQueue:      "test-control-queue",
+			expectTask:        true,
+		},
+		{
+			name:              "no task when feature disabled",
+			featureEnabled:    false,
+			scheduledEventIDs: []int64{5, 6, 7},
+			controlQueue:      "test-control-queue",
+			expectTask:        false,
+		},
+		{
+			name:              "no task when scheduledEventIDs empty",
+			featureEnabled:    true,
+			scheduledEventIDs: []int64{},
+			controlQueue:      "test-control-queue",
+			expectTask:        false,
+		},
+		{
+			name:              "no task when controlQueue empty",
+			featureEnabled:    true,
+			scheduledEventIDs: []int64{5, 6, 7},
+			controlQueue:      "",
+			expectTask:        false,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			ctrl := gomock.NewController(t)
+
+			mutableState := historyi.NewMockMutableState(ctrl)
+			mutableState.EXPECT().GetWorkflowKey().Return(definition.NewWorkflowKey(
+				tests.NamespaceID.String(), tests.WorkflowID, tests.RunID,
+			)).AnyTimes()
+
+			var capturedTasks []tasks.Task
+			if tc.expectTask {
+				mutableState.EXPECT().AddTasks(gomock.Any()).Do(func(ts ...tasks.Task) {
+					capturedTasks = append(capturedTasks, ts...)
+				}).Times(1)
+			}
+
+			cfg := &configs.Config{
+				EnableActivityCancellationNexusTask: func() bool { return tc.featureEnabled },
+			}
+
+			taskGenerator := NewTaskGenerator(nil, mutableState, cfg, nil, log.NewTestLogger())
+			err := taskGenerator.GenerateNotifyActivityTasks(tc.scheduledEventIDs, tc.controlQueue, enumsspb.ACTIVITY_NOTIFICATION_TYPE_CANCEL)
+			require.NoError(t, err)
+
+			if tc.expectTask {
+				require.Len(t, capturedTasks, 1)
+				notifyTask, ok := capturedTasks[0].(*tasks.NotifyActivityTask)
+				require.True(t, ok)
+				assert.Equal(t, tc.scheduledEventIDs, notifyTask.ScheduledEventIDs)
+				assert.Equal(t, tc.controlQueue, notifyTask.Destination)
+				assert.Equal(t, tests.NamespaceID.String(), notifyTask.NamespaceID)
+				assert.Equal(t, enumsspb.ACTIVITY_NOTIFICATION_TYPE_CANCEL, notifyTask.NotificationType)
+			}
+		})
+	}
+}