From a205fb26ec2ddf484b7dda5b4ccdf72b8e0cfab9 Mon Sep 17 00:00:00 2001
From: Zijian <Shaddoll@users.noreply.github.com>
Date: Mon, 21 Oct 2024 20:38:47 +0000
Subject: [PATCH] Add unit test for nosql create tasks method

---
 .../nosql/nosql_task_store_test.go            | 137 +++++++++++++++++-
 1 file changed, 136 insertions(+), 1 deletion(-)

diff --git a/common/persistence/nosql/nosql_task_store_test.go b/common/persistence/nosql/nosql_task_store_test.go
index fcb0d351436..495834486f8 100644
--- a/common/persistence/nosql/nosql_task_store_test.go
+++ b/common/persistence/nosql/nosql_task_store_test.go
@@ -42,6 +42,7 @@ const (
 	TestDomainID     = "test-domain-id"
 	TestDomainName   = "test-domain"
 	TestTaskListName = "test-tasklist"
+	TestTaskType     = persistence.TaskListTypeDecision
 	TestWorkflowID   = "test-workflow-id"
 	TestRunID        = "test-run-id"
 )
@@ -69,7 +70,7 @@ func setupNoSQLStoreMocks(t *testing.T) (*nosqlTaskStore, *nosqlplugin.MockDB) {
 		GetStoreShardByTaskList(
 			TestDomainID,
 			TestTaskListName,
-			int(types.TaskListTypeDecision)).
+			TestTaskType).
 		Return(&nosqlSt, nil).
 		AnyTimes()
 
@@ -444,6 +445,140 @@ func TestCompleteTasksLessThan(t *testing.T) {
 	assert.Equal(t, 13, resp.TasksCompleted)
 }
 
+func TestCreateTasks(t *testing.T) {
+	testCases := []struct {
+		name          string
+		setupMock     func(*nosqlplugin.MockDB)
+		request       *persistence.InternalCreateTasksRequest
+		expectError   bool
+		expectedError string
+	}{
+		{
+			name: "success",
+			setupMock: func(dbMock *nosqlplugin.MockDB) {
+				dbMock.EXPECT().InsertTasks(gomock.Any(), gomock.Any(), &nosqlplugin.TaskListRow{
+					DomainID:     TestDomainID,
+					TaskListName: TestTaskListName,
+					TaskListType: TestTaskType,
+					RangeID:      1,
+				}).Do(func(_ context.Context, tasks []*nosqlplugin.TaskRowForInsert, _ *nosqlplugin.TaskListRow) {
+					assert.Len(t, tasks, 1)
+					assert.Equal(t, TestDomainID, tasks[0].DomainID)
+					assert.Equal(t, "workflow1", tasks[0].WorkflowID)
+					assert.Equal(t, "run1", tasks[0].RunID)
+					assert.Equal(t, int64(100), tasks[0].TaskID)
+					assert.Equal(t, int64(10), tasks[0].ScheduledID)
+					assert.Equal(t, TestTaskType, tasks[0].TaskListType)
+					assert.Equal(t, TestTaskListName, tasks[0].TaskListName)
+					assert.Equal(t, 30, tasks[0].TTLSeconds)
+				}).Return(nil).Times(1)
+			},
+			request: &persistence.InternalCreateTasksRequest{
+				TaskListInfo: &persistence.TaskListInfo{
+					DomainID: TestDomainID,
+					Name:     TestTaskListName,
+					TaskType: TestTaskType,
+					RangeID:  1,
+				},
+				Tasks: []*persistence.InternalCreateTasksInfo{
+					{
+						TaskID: 100,
+						Data: &persistence.InternalTaskInfo{
+							WorkflowID:             "workflow1",
+							RunID:                  "run1",
+							ScheduleID:             10,
+							PartitionConfig:        nil,
+							ScheduleToStartTimeout: 30 * time.Second,
+						},
+					},
+				},
+			},
+			expectError: false,
+		},
+		{
+			name: "condition failure",
+			setupMock: func(dbMock *nosqlplugin.MockDB) {
+				dbMock.EXPECT().InsertTasks(gomock.Any(), gomock.Any(), gomock.Any()).Return(&nosqlplugin.TaskOperationConditionFailure{
+					Details: "rangeID mismatch",
+				}).Times(1)
+			},
+			request: &persistence.InternalCreateTasksRequest{
+				TaskListInfo: &persistence.TaskListInfo{
+					DomainID: TestDomainID,
+					Name:     TestTaskListName,
+					TaskType: TestTaskType,
+					RangeID:  1,
+				},
+				Tasks: []*persistence.InternalCreateTasksInfo{
+					{
+						TaskID: 100,
+						Data: &persistence.InternalTaskInfo{
+							WorkflowID:             "workflow1",
+							RunID:                  "run1",
+							ScheduleID:             10,
+							PartitionConfig:        nil,
+							ScheduleToStartTimeout: 30 * time.Second,
+						},
+					},
+				},
+			},
+			expectError:   true,
+			expectedError: "Failed to insert tasks. name: test-tasklist, type: 0, rangeID: 1, columns: (rangeID mismatch)",
+		},
+		{
+			name: "generic db error",
+			setupMock: func(dbMock *nosqlplugin.MockDB) {
+				dbMock.EXPECT().InsertTasks(gomock.Any(), gomock.Any(), gomock.Any()).Return(errors.New("db error")).Times(1)
+				dbMock.EXPECT().IsNotFoundError(gomock.Any()).Return(true).Times(1)
+			},
+			request: &persistence.InternalCreateTasksRequest{
+				TaskListInfo: &persistence.TaskListInfo{
+					DomainID: TestDomainID,
+					Name:     TestTaskListName,
+					TaskType: TestTaskType,
+					RangeID:  1,
+				},
+				Tasks: []*persistence.InternalCreateTasksInfo{
+					{
+						TaskID: 100,
+						Data: &persistence.InternalTaskInfo{
+							WorkflowID:             "workflow1",
+							RunID:                  "run1",
+							ScheduleID:             10,
+							PartitionConfig:        nil,
+							ScheduleToStartTimeout: 30 * time.Second,
+						},
+					},
+				},
+			},
+			expectError:   true,
+			expectedError: "CreateTasks failed. Error: db error",
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			// Setup mocks for each test case individually
+			store, dbMock := setupNoSQLStoreMocks(t)
+
+			// Setup test-specific mock behavior
+			tc.setupMock(dbMock)
+
+			// Execute the method under test
+			resp, err := store.CreateTasks(context.Background(), tc.request)
+
+			// Validate results
+			if tc.expectError {
+				assert.Error(t, err)
+				assert.Contains(t, err.Error(), tc.expectedError)
+			} else {
+				assert.NoError(t, err)
+				assert.NotNil(t, resp)
+			}
+		})
+	}
+}
+
 func getValidLeaseTaskListRequest() *persistence.LeaseTaskListRequest {
 	return &persistence.LeaseTaskListRequest{
 		DomainID:     TestDomainID,