diff --git a/common/persistence/pinot/pinot_visibility_store_test.go b/common/persistence/pinot/pinot_visibility_store_test.go
index 3aced82094f..1d7abddc8a8 100644
--- a/common/persistence/pinot/pinot_visibility_store_test.go
+++ b/common/persistence/pinot/pinot_visibility_store_test.go
@@ -89,19 +89,37 @@ func TestRecordWorkflowExecutionStarted(t *testing.T) {
 	}
 
 	tests := map[string]struct {
-		request       *p.InternalRecordWorkflowExecutionStartedRequest
-		expectedError error
+		request                *p.InternalRecordWorkflowExecutionStartedRequest
+		producerMockAffordance func(mockProducer *mocks.KafkaProducer)
+		expectedError          error
 	}{
 		"Case1: error case": {
-			request:       errorRequest,
-			expectedError: fmt.Errorf("error"),
+			request: errorRequest,
+			producerMockAffordance: func(mockProducer *mocks.KafkaProducer) {
+				mockProducer.On("Publish", mock.Anything, mock.MatchedBy(func(input *indexer.PinotMessage) bool {
+					return true
+				})).Return(&types.BadRequestError{}).Once()
+			},
+			expectedError: fmt.Errorf("invalid character 'e' in literal true (expecting 'r')"),
 		},
 		"Case2: normal case": {
-			request:       request,
+			request: request,
+			producerMockAffordance: func(mockProducer *mocks.KafkaProducer) {
+				mockProducer.On("Publish", mock.Anything, mock.MatchedBy(func(input *indexer.PinotMessage) bool {
+					assert.Equal(t, request.WorkflowID, input.GetWorkflowID())
+					return true
+				})).Return(nil).Once()
+			},
 			expectedError: nil,
 		},
 		"Case3: normal case with search attributes": {
-			request:       requestWithSearchAttributes,
+			request: requestWithSearchAttributes,
+			producerMockAffordance: func(mockProducer *mocks.KafkaProducer) {
+				mockProducer.On("Publish", mock.Anything, mock.MatchedBy(func(input *indexer.PinotMessage) bool {
+					assert.Equal(t, request.WorkflowID, input.GetWorkflowID())
+					return true
+				})).Return(nil).Once()
+			},
 			expectedError: nil,
 		},
 	}
@@ -117,14 +135,11 @@ func TestRecordWorkflowExecutionStarted(t *testing.T) {
 			}, mockProducer, log.NewNoop())
 			visibilityStore := mgr.(*pinotVisibilityStore)
 
-			mockProducer.On("Publish", mock.Anything, mock.MatchedBy(func(input *indexer.PinotMessage) bool {
-				assert.Equal(t, request.WorkflowID, input.GetWorkflowID())
-				return true
-			})).Return(nil).Once()
+			test.producerMockAffordance(mockProducer)
 
 			err := visibilityStore.RecordWorkflowExecutionStarted(context.Background(), test.request)
 			if test.expectedError != nil {
-				assert.Error(t, err)
+				assert.Equal(t, test.expectedError.Error(), err.Error())
 			} else {
 				assert.NoError(t, err)
 			}
@@ -135,7 +150,7 @@ func TestRecordWorkflowExecutionStarted(t *testing.T) {
 func TestRecordWorkflowExecutionClosed(t *testing.T) {
 	// test non-empty request fields match
 	errorRequest := &p.InternalRecordWorkflowExecutionClosedRequest{
-		WorkflowID: "wid",
+		WorkflowID: "error-wid",
 		Memo:       p.NewDataBlob([]byte(`test bytes`), common.EncodingTypeThriftRW),
 		SearchAttributes: map[string][]byte{
 			"CustomStringField": []byte("test string"),
@@ -148,15 +163,27 @@ func TestRecordWorkflowExecutionClosed(t *testing.T) {
 	}
 
 	tests := map[string]struct {
-		request       *p.InternalRecordWorkflowExecutionClosedRequest
-		expectedError error
+		request                *p.InternalRecordWorkflowExecutionClosedRequest
+		producerMockAffordance func(mockProducer *mocks.KafkaProducer)
+		expectedError          error
 	}{
 		"Case1: error case": {
-			request:       errorRequest,
-			expectedError: fmt.Errorf("error"),
+			request: errorRequest,
+			producerMockAffordance: func(mockProducer *mocks.KafkaProducer) {
+				mockProducer.On("Publish", mock.Anything, mock.MatchedBy(func(input *indexer.PinotMessage) bool {
+					return true
+				})).Return(&types.BadRequestError{}).Once()
+			},
+			expectedError: fmt.Errorf("invalid character 'e' in literal true (expecting 'r')"),
 		},
 		"Case2: normal case": {
-			request:       request,
+			request: request,
+			producerMockAffordance: func(mockProducer *mocks.KafkaProducer) {
+				mockProducer.On("Publish", mock.Anything, mock.MatchedBy(func(input *indexer.PinotMessage) bool {
+					assert.Equal(t, request.WorkflowID, input.GetWorkflowID())
+					return true
+				})).Return(nil).Once()
+			},
 			expectedError: nil,
 		},
 	}
@@ -172,14 +199,11 @@ func TestRecordWorkflowExecutionClosed(t *testing.T) {
 			}, mockProducer, log.NewNoop())
 			visibilityStore := mgr.(*pinotVisibilityStore)
 
-			mockProducer.On("Publish", mock.Anything, mock.MatchedBy(func(input *indexer.PinotMessage) bool {
-				assert.Equal(t, request.WorkflowID, input.GetWorkflowID())
-				return true
-			})).Return(nil).Once()
+			test.producerMockAffordance(mockProducer)
 
 			err := visibilityStore.RecordWorkflowExecutionClosed(context.Background(), test.request)
 			if test.expectedError != nil {
-				assert.Error(t, err)
+				assert.Equal(t, test.expectedError.Error(), err.Error())
 			} else {
 				assert.NoError(t, err)
 			}
@@ -188,17 +212,23 @@ func TestRecordWorkflowExecutionClosed(t *testing.T) {
 }
 
 func TestRecordWorkflowExecutionUninitialized(t *testing.T) {
-	// test non-empty request fields match
 	request := &p.InternalRecordWorkflowExecutionUninitializedRequest{
 		WorkflowID: "wid",
 	}
 
 	tests := map[string]struct {
-		request       *p.InternalRecordWorkflowExecutionUninitializedRequest
-		expectedError error
+		request                *p.InternalRecordWorkflowExecutionUninitializedRequest
+		producerMockAffordance func(mockProducer *mocks.KafkaProducer)
+		expectedError          error
 	}{
 		"Case1: normal case": {
-			request:       request,
+			request: request,
+			producerMockAffordance: func(mockProducer *mocks.KafkaProducer) {
+				mockProducer.On("Publish", mock.Anything, mock.MatchedBy(func(input *indexer.PinotMessage) bool {
+					assert.Equal(t, request.WorkflowID, input.GetWorkflowID())
+					return true
+				})).Return(nil).Once()
+			},
 			expectedError: nil,
 		},
 	}
@@ -214,10 +244,7 @@ func TestRecordWorkflowExecutionUninitialized(t *testing.T) {
 			}, mockProducer, log.NewNoop())
 			visibilityStore := mgr.(*pinotVisibilityStore)
 
-			mockProducer.On("Publish", mock.Anything, mock.MatchedBy(func(input *indexer.PinotMessage) bool {
-				assert.Equal(t, request.WorkflowID, input.GetWorkflowID())
-				return true
-			})).Return(nil).Once()
+			test.producerMockAffordance(mockProducer)
 
 			err := visibilityStore.RecordWorkflowExecutionUninitialized(context.Background(), test.request)
 			assert.Equal(t, test.expectedError, err)
@@ -228,7 +255,7 @@ func TestRecordWorkflowExecutionUninitialized(t *testing.T) {
 func TestUpsertWorkflowExecution(t *testing.T) {
 	// test non-empty request fields match
 	errorRequest := &p.InternalUpsertWorkflowExecutionRequest{
-		WorkflowID: "wid",
+		WorkflowID: "error-wid",
 		Memo:       p.NewDataBlob([]byte(`test bytes`), common.EncodingTypeThriftRW),
 		SearchAttributes: map[string][]byte{
 			"CustomStringField": []byte("test string"),
@@ -241,15 +268,27 @@ func TestUpsertWorkflowExecution(t *testing.T) {
 	request.Memo = p.NewDataBlob(memoBytes, common.EncodingTypeThriftRW)
 
 	tests := map[string]struct {
-		request       *p.InternalUpsertWorkflowExecutionRequest
-		expectedError error
+		request                *p.InternalUpsertWorkflowExecutionRequest
+		producerMockAffordance func(mockProducer *mocks.KafkaProducer)
+		expectedError          error
 	}{
 		"Case1: error case": {
-			request:       errorRequest,
-			expectedError: fmt.Errorf("error"),
+			request: errorRequest,
+			producerMockAffordance: func(mockProducer *mocks.KafkaProducer) {
+				mockProducer.On("Publish", mock.Anything, mock.MatchedBy(func(input *indexer.PinotMessage) bool {
+					return true
+				})).Return(fmt.Errorf("error")).Once()
+			},
+			expectedError: fmt.Errorf("invalid character 'e' in literal true (expecting 'r')"),
 		},
 		"Case2: normal case": {
-			request:       request,
+			request: request,
+			producerMockAffordance: func(mockProducer *mocks.KafkaProducer) {
+				mockProducer.On("Publish", mock.Anything, mock.MatchedBy(func(input *indexer.PinotMessage) bool {
+					assert.Equal(t, request.WorkflowID, input.GetWorkflowID())
+					return true
+				})).Return(nil).Once()
+			},
 			expectedError: nil,
 		},
 	}
@@ -265,14 +304,11 @@ func TestUpsertWorkflowExecution(t *testing.T) {
 			}, mockProducer, log.NewNoop())
 			visibilityStore := mgr.(*pinotVisibilityStore)
 
-			mockProducer.On("Publish", mock.Anything, mock.MatchedBy(func(input *indexer.PinotMessage) bool {
-				assert.Equal(t, request.WorkflowID, input.GetWorkflowID())
-				return true
-			})).Return(nil).Once()
+			test.producerMockAffordance(mockProducer)
 
 			err := visibilityStore.UpsertWorkflowExecution(context.Background(), test.request)
 			if test.expectedError != nil {
-				assert.Error(t, err)
+				assert.Equal(t, test.expectedError.Error(), err.Error())
 			} else {
 				assert.NoError(t, err)
 			}
@@ -286,11 +322,18 @@ func TestDeleteWorkflowExecution(t *testing.T) {
 	request.WorkflowID = "wid"
 
 	tests := map[string]struct {
-		request       *p.VisibilityDeleteWorkflowExecutionRequest
-		expectedError error
+		request                *p.VisibilityDeleteWorkflowExecutionRequest
+		producerMockAffordance func(mockProducer *mocks.KafkaProducer)
+		expectedError          error
 	}{
 		"Case1: normal case": {
-			request:       request,
+			request: request,
+			producerMockAffordance: func(mockProducer *mocks.KafkaProducer) {
+				mockProducer.On("Publish", mock.Anything, mock.MatchedBy(func(input *indexer.PinotMessage) bool {
+					assert.Equal(t, request.WorkflowID, input.GetWorkflowID())
+					return true
+				})).Return(nil).Once()
+			},
 			expectedError: nil,
 		},
 	}
@@ -306,10 +349,7 @@ func TestDeleteWorkflowExecution(t *testing.T) {
 			}, mockProducer, log.NewNoop())
 			visibilityStore := mgr.(*pinotVisibilityStore)
 
-			mockProducer.On("Publish", mock.Anything, mock.MatchedBy(func(input *indexer.PinotMessage) bool {
-				assert.Equal(t, request.WorkflowID, input.GetWorkflowID())
-				return true
-			})).Return(nil).Once()
+			test.producerMockAffordance(mockProducer)
 
 			err := visibilityStore.DeleteWorkflowExecution(context.Background(), test.request)
 			assert.Equal(t, test.expectedError, err)
@@ -328,11 +368,18 @@ func TestDeleteUninitializedWorkflowExecution(t *testing.T) {
 	}
 
 	tests := map[string]struct {
-		request       *p.VisibilityDeleteWorkflowExecutionRequest
-		expectedError error
+		request                *p.VisibilityDeleteWorkflowExecutionRequest
+		producerMockAffordance func(mockProducer *mocks.KafkaProducer)
+		expectedError          error
 	}{
 		"Case1: normal case": {
-			request:       request,
+			request: request,
+			producerMockAffordance: func(mockProducer *mocks.KafkaProducer) {
+				mockProducer.On("Publish", mock.Anything, mock.MatchedBy(func(input *indexer.PinotMessage) bool {
+					assert.Equal(t, request.WorkflowID, input.GetWorkflowID())
+					return true
+				})).Return(nil).Once()
+			},
 			expectedError: nil,
 		},
 	}
@@ -351,10 +398,7 @@ func TestDeleteUninitializedWorkflowExecution(t *testing.T) {
 			mockPinotClient.EXPECT().GetTableName().Return(testTableName).Times(1)
 			mockPinotClient.EXPECT().CountByQuery(gomock.Any()).Return(int64(1), nil).Times(1)
 
-			mockProducer.On("Publish", mock.Anything, mock.MatchedBy(func(input *indexer.PinotMessage) bool {
-				assert.Equal(t, request.WorkflowID, input.GetWorkflowID())
-				return true
-			})).Return(nil).Once()
+			test.producerMockAffordance(mockProducer)
 
 			err := visibilityStore.DeleteUninitializedWorkflowExecution(context.Background(), test.request)
 			assert.Equal(t, test.expectedError, err)
@@ -373,18 +417,26 @@ func TestListOpenWorkflowExecutions(t *testing.T) {
 	}
 
 	tests := map[string]struct {
-		request       *p.InternalListWorkflowExecutionsRequest
-		expectedResp  *p.InternalListWorkflowExecutionsResponse
-		expectedError error
+		request                   *p.InternalListWorkflowExecutionsRequest
+		expectedResp              *p.InternalListWorkflowExecutionsResponse
+		pinotClientMockAffordance func(mockPinotClient *pnt.MockGenericClient)
+		expectedError             error
 	}{
 		"Case1: error case": {
-			request:       errorRequest,
-			expectedResp:  nil,
-			expectedError: fmt.Errorf("error"),
+			request:      errorRequest,
+			expectedResp: nil,
+			pinotClientMockAffordance: func(mockPinotClient *pnt.MockGenericClient) {
+				mockPinotClient.EXPECT().GetTableName().Return(testTableName).Times(1)
+			},
+			expectedError: fmt.Errorf("next page token: unable to deserialize page token. err: invalid character 'e' looking for beginning of value"),
 		},
 		"Case2: normal case with nil response": {
-			request:       request,
-			expectedResp:  nil,
+			request:      request,
+			expectedResp: nil,
+			pinotClientMockAffordance: func(mockPinotClient *pnt.MockGenericClient) {
+				mockPinotClient.EXPECT().GetTableName().Return(testTableName).Times(1)
+				mockPinotClient.EXPECT().Search(gomock.Any()).Return(nil, nil).Times(1)
+			},
 			expectedError: nil,
 		},
 	}
@@ -400,16 +452,13 @@ func TestListOpenWorkflowExecutions(t *testing.T) {
 			}, mockProducer, log.NewNoop())
 			visibilityStore := mgr.(*pinotVisibilityStore)
 
+			test.pinotClientMockAffordance(mockPinotClient)
+
+			resp, err := visibilityStore.ListOpenWorkflowExecutions(context.Background(), test.request)
+			assert.Equal(t, test.expectedResp, resp)
 			if test.expectedError != nil {
-				mockPinotClient.EXPECT().GetTableName().Return(testTableName).Times(1)
-				resp, err := visibilityStore.ListOpenWorkflowExecutions(context.Background(), test.request)
-				assert.Equal(t, test.expectedResp, resp)
-				assert.Error(t, err)
+				assert.Equal(t, test.expectedError.Error(), err.Error())
 			} else {
-				mockPinotClient.EXPECT().GetTableName().Return(testTableName).Times(1)
-				mockPinotClient.EXPECT().Search(gomock.Any()).Return(nil, nil).Times(1)
-				resp, err := visibilityStore.ListOpenWorkflowExecutions(context.Background(), test.request)
-				assert.Equal(t, test.expectedResp, resp)
 				assert.NoError(t, err)
 			}
 		})
@@ -427,18 +476,26 @@ func TestListClosedWorkflowExecutions(t *testing.T) {
 	}
 
 	tests := map[string]struct {
-		request       *p.InternalListWorkflowExecutionsRequest
-		expectedResp  *p.InternalListWorkflowExecutionsResponse
-		expectedError error
+		request                   *p.InternalListWorkflowExecutionsRequest
+		expectedResp              *p.InternalListWorkflowExecutionsResponse
+		pinotClientMockAffordance func(mockPinotClient *pnt.MockGenericClient)
+		expectedError             error
 	}{
 		"Case1: error case": {
-			request:       errorRequest,
-			expectedResp:  nil,
-			expectedError: fmt.Errorf("error"),
+			request:      errorRequest,
+			expectedResp: nil,
+			pinotClientMockAffordance: func(mockPinotClient *pnt.MockGenericClient) {
+				mockPinotClient.EXPECT().GetTableName().Return(testTableName).Times(1)
+			},
+			expectedError: fmt.Errorf("next page token: unable to deserialize page token. err: invalid character 'e' looking for beginning of value"),
 		},
 		"Case2: normal case with nil response": {
-			request:       request,
-			expectedResp:  nil,
+			request:      request,
+			expectedResp: nil,
+			pinotClientMockAffordance: func(mockPinotClient *pnt.MockGenericClient) {
+				mockPinotClient.EXPECT().GetTableName().Return(testTableName).Times(1)
+				mockPinotClient.EXPECT().Search(gomock.Any()).Return(nil, nil).Times(1)
+			},
 			expectedError: nil,
 		},
 	}
@@ -454,16 +511,13 @@ func TestListClosedWorkflowExecutions(t *testing.T) {
 			}, mockProducer, log.NewNoop())
 			visibilityStore := mgr.(*pinotVisibilityStore)
 
+			test.pinotClientMockAffordance(mockPinotClient)
+
+			resp, err := visibilityStore.ListClosedWorkflowExecutions(context.Background(), test.request)
+			assert.Equal(t, test.expectedResp, resp)
 			if test.expectedError != nil {
-				mockPinotClient.EXPECT().GetTableName().Return(testTableName).Times(1)
-				resp, err := visibilityStore.ListClosedWorkflowExecutions(context.Background(), test.request)
-				assert.Equal(t, test.expectedResp, resp)
-				assert.Error(t, err)
+				assert.Equal(t, test.expectedError.Error(), err.Error())
 			} else {
-				mockPinotClient.EXPECT().GetTableName().Return(testTableName).Times(1)
-				mockPinotClient.EXPECT().Search(gomock.Any()).Return(nil, nil).Times(1)
-				resp, err := visibilityStore.ListClosedWorkflowExecutions(context.Background(), test.request)
-				assert.Equal(t, test.expectedResp, resp)
 				assert.NoError(t, err)
 			}
 		})
@@ -482,18 +536,26 @@ func TestListOpenWorkflowExecutionsByType(t *testing.T) {
 	request := &p.InternalListWorkflowExecutionsByTypeRequest{}
 
 	tests := map[string]struct {
-		request       *p.InternalListWorkflowExecutionsByTypeRequest
-		expectedResp  *p.InternalListWorkflowExecutionsResponse
-		expectedError error
+		request                   *p.InternalListWorkflowExecutionsByTypeRequest
+		expectedResp              *p.InternalListWorkflowExecutionsResponse
+		pinotClientMockAffordance func(mockPinotClient *pnt.MockGenericClient)
+		expectedError             error
 	}{
 		"Case1: error case": {
-			request:       errorRequest,
-			expectedResp:  nil,
-			expectedError: fmt.Errorf("error"),
+			request:      errorRequest,
+			expectedResp: nil,
+			pinotClientMockAffordance: func(mockPinotClient *pnt.MockGenericClient) {
+				mockPinotClient.EXPECT().GetTableName().Return(testTableName).Times(1)
+			},
+			expectedError: fmt.Errorf("next page token: unable to deserialize page token. err: invalid character 'e' looking for beginning of value"),
 		},
 		"Case2: normal case with nil response": {
-			request:       request,
-			expectedResp:  nil,
+			request:      request,
+			expectedResp: nil,
+			pinotClientMockAffordance: func(mockPinotClient *pnt.MockGenericClient) {
+				mockPinotClient.EXPECT().GetTableName().Return(testTableName).Times(1)
+				mockPinotClient.EXPECT().Search(gomock.Any()).Return(nil, nil).Times(1)
+			},
 			expectedError: nil,
 		},
 	}
@@ -509,16 +571,13 @@ func TestListOpenWorkflowExecutionsByType(t *testing.T) {
 			}, mockProducer, log.NewNoop())
 			visibilityStore := mgr.(*pinotVisibilityStore)
 
+			test.pinotClientMockAffordance(mockPinotClient)
+
+			resp, err := visibilityStore.ListOpenWorkflowExecutionsByType(context.Background(), test.request)
+			assert.Equal(t, test.expectedResp, resp)
 			if test.expectedError != nil {
-				mockPinotClient.EXPECT().GetTableName().Return(testTableName).Times(1)
-				resp, err := visibilityStore.ListOpenWorkflowExecutionsByType(context.Background(), test.request)
-				assert.Equal(t, test.expectedResp, resp)
-				assert.Error(t, err)
+				assert.Equal(t, test.expectedError.Error(), err.Error())
 			} else {
-				mockPinotClient.EXPECT().GetTableName().Return(testTableName).Times(1)
-				mockPinotClient.EXPECT().Search(gomock.Any()).Return(nil, nil).Times(1)
-				resp, err := visibilityStore.ListOpenWorkflowExecutionsByType(context.Background(), test.request)
-				assert.Equal(t, test.expectedResp, resp)
 				assert.NoError(t, err)
 			}
 		})
@@ -536,18 +595,26 @@ func TestListClosedWorkflowExecutionsByType(t *testing.T) {
 	request := &p.InternalListWorkflowExecutionsByTypeRequest{}
 
 	tests := map[string]struct {
-		request       *p.InternalListWorkflowExecutionsByTypeRequest
-		expectedResp  *p.InternalListWorkflowExecutionsResponse
-		expectedError error
+		request                   *p.InternalListWorkflowExecutionsByTypeRequest
+		expectedResp              *p.InternalListWorkflowExecutionsResponse
+		pinotClientMockAffordance func(mockPinotClient *pnt.MockGenericClient)
+		expectedError             error
 	}{
 		"Case1: error case": {
-			request:       errorRequest,
-			expectedResp:  nil,
-			expectedError: fmt.Errorf("error"),
+			request:      errorRequest,
+			expectedResp: nil,
+			pinotClientMockAffordance: func(mockPinotClient *pnt.MockGenericClient) {
+				mockPinotClient.EXPECT().GetTableName().Return(testTableName).Times(1)
+			},
+			expectedError: fmt.Errorf("next page token: unable to deserialize page token. err: invalid character 'e' looking for beginning of value"),
 		},
 		"Case2: normal case with nil response": {
-			request:       request,
-			expectedResp:  nil,
+			request:      request,
+			expectedResp: nil,
+			pinotClientMockAffordance: func(mockPinotClient *pnt.MockGenericClient) {
+				mockPinotClient.EXPECT().GetTableName().Return(testTableName).Times(1)
+				mockPinotClient.EXPECT().Search(gomock.Any()).Return(nil, nil).Times(1)
+			},
 			expectedError: nil,
 		},
 	}
@@ -563,16 +630,13 @@ func TestListClosedWorkflowExecutionsByType(t *testing.T) {
 			}, mockProducer, log.NewNoop())
 			visibilityStore := mgr.(*pinotVisibilityStore)
 
+			test.pinotClientMockAffordance(mockPinotClient)
+
+			resp, err := visibilityStore.ListClosedWorkflowExecutionsByType(context.Background(), test.request)
+			assert.Equal(t, test.expectedResp, resp)
 			if test.expectedError != nil {
-				mockPinotClient.EXPECT().GetTableName().Return(testTableName).Times(1)
-				resp, err := visibilityStore.ListClosedWorkflowExecutionsByType(context.Background(), test.request)
-				assert.Equal(t, test.expectedResp, resp)
-				assert.Error(t, err)
+				assert.Equal(t, test.expectedError.Error(), err.Error())
 			} else {
-				mockPinotClient.EXPECT().GetTableName().Return(testTableName).Times(1)
-				mockPinotClient.EXPECT().Search(gomock.Any()).Return(nil, nil).Times(1)
-				resp, err := visibilityStore.ListClosedWorkflowExecutionsByType(context.Background(), test.request)
-				assert.Equal(t, test.expectedResp, resp)
 				assert.NoError(t, err)
 			}
 		})
@@ -589,18 +653,26 @@ func TestListOpenWorkflowExecutionsByWorkflowID(t *testing.T) {
 	request := &p.InternalListWorkflowExecutionsByWorkflowIDRequest{}
 
 	tests := map[string]struct {
-		request       *p.InternalListWorkflowExecutionsByWorkflowIDRequest
-		expectedResp  *p.InternalListWorkflowExecutionsResponse
-		expectedError error
+		request                   *p.InternalListWorkflowExecutionsByWorkflowIDRequest
+		expectedResp              *p.InternalListWorkflowExecutionsResponse
+		pinotClientMockAffordance func(mockPinotClient *pnt.MockGenericClient)
+		expectedError             error
 	}{
 		"Case1: error case": {
-			request:       errorRequest,
-			expectedResp:  nil,
-			expectedError: fmt.Errorf("error"),
+			request:      errorRequest,
+			expectedResp: nil,
+			pinotClientMockAffordance: func(mockPinotClient *pnt.MockGenericClient) {
+				mockPinotClient.EXPECT().GetTableName().Return(testTableName).Times(1)
+			},
+			expectedError: fmt.Errorf("next page token: unable to deserialize page token. err: invalid character 'e' looking for beginning of value"),
 		},
 		"Case2: normal case with nil response": {
-			request:       request,
-			expectedResp:  nil,
+			request:      request,
+			expectedResp: nil,
+			pinotClientMockAffordance: func(mockPinotClient *pnt.MockGenericClient) {
+				mockPinotClient.EXPECT().GetTableName().Return(testTableName).Times(1)
+				mockPinotClient.EXPECT().Search(gomock.Any()).Return(nil, nil).Times(1)
+			},
 			expectedError: nil,
 		},
 	}
@@ -616,16 +688,13 @@ func TestListOpenWorkflowExecutionsByWorkflowID(t *testing.T) {
 			}, mockProducer, log.NewNoop())
 			visibilityStore := mgr.(*pinotVisibilityStore)
 
+			test.pinotClientMockAffordance(mockPinotClient)
+
+			resp, err := visibilityStore.ListOpenWorkflowExecutionsByWorkflowID(context.Background(), test.request)
+			assert.Equal(t, test.expectedResp, resp)
 			if test.expectedError != nil {
-				mockPinotClient.EXPECT().GetTableName().Return(testTableName).Times(1)
-				resp, err := visibilityStore.ListOpenWorkflowExecutionsByWorkflowID(context.Background(), test.request)
-				assert.Equal(t, test.expectedResp, resp)
-				assert.Error(t, err)
+				assert.Equal(t, test.expectedError.Error(), err.Error())
 			} else {
-				mockPinotClient.EXPECT().GetTableName().Return(testTableName).Times(1)
-				mockPinotClient.EXPECT().Search(gomock.Any()).Return(nil, nil).Times(1)
-				resp, err := visibilityStore.ListOpenWorkflowExecutionsByWorkflowID(context.Background(), test.request)
-				assert.Equal(t, test.expectedResp, resp)
 				assert.NoError(t, err)
 			}
 		})
@@ -642,18 +711,26 @@ func TestListClosedWorkflowExecutionsByWorkflowID(t *testing.T) {
 	request := &p.InternalListWorkflowExecutionsByWorkflowIDRequest{}
 
 	tests := map[string]struct {
-		request       *p.InternalListWorkflowExecutionsByWorkflowIDRequest
-		expectedResp  *p.InternalListWorkflowExecutionsResponse
-		expectedError error
+		request                   *p.InternalListWorkflowExecutionsByWorkflowIDRequest
+		expectedResp              *p.InternalListWorkflowExecutionsResponse
+		pinotClientMockAffordance func(mockPinotClient *pnt.MockGenericClient)
+		expectedError             error
 	}{
 		"Case1: error case": {
-			request:       errorRequest,
-			expectedResp:  nil,
-			expectedError: fmt.Errorf("error"),
+			request:      errorRequest,
+			expectedResp: nil,
+			pinotClientMockAffordance: func(mockPinotClient *pnt.MockGenericClient) {
+				mockPinotClient.EXPECT().GetTableName().Return(testTableName).Times(1)
+			},
+			expectedError: fmt.Errorf("next page token: unable to deserialize page token. err: invalid character 'e' looking for beginning of value"),
 		},
 		"Case2: normal case with nil response": {
-			request:       request,
-			expectedResp:  nil,
+			request:      request,
+			expectedResp: nil,
+			pinotClientMockAffordance: func(mockPinotClient *pnt.MockGenericClient) {
+				mockPinotClient.EXPECT().GetTableName().Return(testTableName).Times(1)
+				mockPinotClient.EXPECT().Search(gomock.Any()).Return(nil, nil).Times(1)
+			},
 			expectedError: nil,
 		},
 	}
@@ -669,16 +746,13 @@ func TestListClosedWorkflowExecutionsByWorkflowID(t *testing.T) {
 			}, mockProducer, log.NewNoop())
 			visibilityStore := mgr.(*pinotVisibilityStore)
 
+			test.pinotClientMockAffordance(mockPinotClient)
+
+			resp, err := visibilityStore.ListClosedWorkflowExecutionsByWorkflowID(context.Background(), test.request)
+			assert.Equal(t, test.expectedResp, resp)
 			if test.expectedError != nil {
-				mockPinotClient.EXPECT().GetTableName().Return(testTableName).Times(1)
-				resp, err := visibilityStore.ListClosedWorkflowExecutionsByWorkflowID(context.Background(), test.request)
-				assert.Equal(t, test.expectedResp, resp)
-				assert.Error(t, err)
+				assert.Equal(t, test.expectedError.Error(), err.Error())
 			} else {
-				mockPinotClient.EXPECT().GetTableName().Return(testTableName).Times(1)
-				mockPinotClient.EXPECT().Search(gomock.Any()).Return(nil, nil).Times(1)
-				resp, err := visibilityStore.ListClosedWorkflowExecutionsByWorkflowID(context.Background(), test.request)
-				assert.Equal(t, test.expectedResp, resp)
 				assert.NoError(t, err)
 			}
 		})
@@ -695,18 +769,26 @@ func TestListClosedWorkflowExecutionsByStatus(t *testing.T) {
 	request := &p.InternalListClosedWorkflowExecutionsByStatusRequest{}
 
 	tests := map[string]struct {
-		request       *p.InternalListClosedWorkflowExecutionsByStatusRequest
-		expectedResp  *p.InternalListWorkflowExecutionsResponse
-		expectedError error
+		request                   *p.InternalListClosedWorkflowExecutionsByStatusRequest
+		expectedResp              *p.InternalListWorkflowExecutionsResponse
+		pinotClientMockAffordance func(mockPinotClient *pnt.MockGenericClient)
+		expectedError             error
 	}{
 		"Case1: error case": {
-			request:       errorRequest,
-			expectedResp:  nil,
-			expectedError: fmt.Errorf("error"),
+			request:      errorRequest,
+			expectedResp: nil,
+			pinotClientMockAffordance: func(mockPinotClient *pnt.MockGenericClient) {
+				mockPinotClient.EXPECT().GetTableName().Return(testTableName).Times(1)
+			},
+			expectedError: fmt.Errorf("next page token: unable to deserialize page token. err: invalid character 'e' looking for beginning of value"),
 		},
 		"Case2: normal case with nil response": {
-			request:       request,
-			expectedResp:  nil,
+			request:      request,
+			expectedResp: nil,
+			pinotClientMockAffordance: func(mockPinotClient *pnt.MockGenericClient) {
+				mockPinotClient.EXPECT().GetTableName().Return(testTableName).Times(1)
+				mockPinotClient.EXPECT().Search(gomock.Any()).Return(nil, nil).Times(1)
+			},
 			expectedError: nil,
 		},
 	}
@@ -722,16 +804,13 @@ func TestListClosedWorkflowExecutionsByStatus(t *testing.T) {
 			}, mockProducer, log.NewNoop())
 			visibilityStore := mgr.(*pinotVisibilityStore)
 
+			test.pinotClientMockAffordance(mockPinotClient)
+
+			resp, err := visibilityStore.ListClosedWorkflowExecutionsByStatus(context.Background(), test.request)
+			assert.Equal(t, test.expectedResp, resp)
 			if test.expectedError != nil {
-				mockPinotClient.EXPECT().GetTableName().Return(testTableName).Times(1)
-				resp, err := visibilityStore.ListClosedWorkflowExecutionsByStatus(context.Background(), test.request)
-				assert.Equal(t, test.expectedResp, resp)
-				assert.Error(t, err)
+				assert.Equal(t, test.expectedError.Error(), err.Error())
 			} else {
-				mockPinotClient.EXPECT().GetTableName().Return(testTableName).Times(1)
-				mockPinotClient.EXPECT().Search(gomock.Any()).Return(nil, nil).Times(1)
-				resp, err := visibilityStore.ListClosedWorkflowExecutionsByStatus(context.Background(), test.request)
-				assert.Equal(t, test.expectedResp, resp)
 				assert.NoError(t, err)
 			}
 		})
@@ -743,17 +822,39 @@ func TestGetClosedWorkflowExecution(t *testing.T) {
 	request := &p.InternalGetClosedWorkflowExecutionRequest{}
 
 	tests := map[string]struct {
-		request       *p.InternalGetClosedWorkflowExecutionRequest
-		expectedResp  *p.InternalGetClosedWorkflowExecutionRequest
-		expectedError error
+		request                   *p.InternalGetClosedWorkflowExecutionRequest
+		expectedResp              *p.InternalGetClosedWorkflowExecutionRequest
+		pinotClientMockAffordance func(mockPinotClient *pnt.MockGenericClient)
+		expectedError             error
 	}{
 		"Case1: error case": {
-			request:       errorRequest,
-			expectedResp:  nil,
-			expectedError: fmt.Errorf("error"),
+			request:      errorRequest,
+			expectedResp: nil,
+			pinotClientMockAffordance: func(mockPinotClient *pnt.MockGenericClient) {
+				mockPinotClient.EXPECT().GetTableName().Return(testTableName).Times(1)
+				mockPinotClient.EXPECT().Search(gomock.Any()).Return(&pnt.SearchResponse{
+					Executions: []*p.InternalVisibilityWorkflowExecutionInfo{
+						{
+							DomainID: DomainID,
+						},
+					},
+				}, fmt.Errorf("error")).Times(1)
+			},
+			expectedError: fmt.Errorf("Pinot GetClosedWorkflowExecution failed, error"),
 		},
 		"Case2: normal case with nil response": {
-			request:       request,
+			request:      request,
+			expectedResp: nil,
+			pinotClientMockAffordance: func(mockPinotClient *pnt.MockGenericClient) {
+				mockPinotClient.EXPECT().GetTableName().Return(testTableName).Times(1)
+				mockPinotClient.EXPECT().Search(gomock.Any()).Return(&pnt.SearchResponse{
+					Executions: []*p.InternalVisibilityWorkflowExecutionInfo{
+						{
+							DomainID: DomainID,
+						},
+					},
+				}, nil).Times(1)
+			},
 			expectedError: nil,
 		},
 	}
@@ -769,27 +870,12 @@ func TestGetClosedWorkflowExecution(t *testing.T) {
 			}, mockProducer, log.NewNoop())
 			visibilityStore := mgr.(*pinotVisibilityStore)
 
+			test.pinotClientMockAffordance(mockPinotClient)
+
+			_, err := visibilityStore.GetClosedWorkflowExecution(context.Background(), test.request)
 			if test.expectedError != nil {
-				mockPinotClient.EXPECT().GetTableName().Return(testTableName).Times(1)
-				mockPinotClient.EXPECT().Search(gomock.Any()).Return(&pnt.SearchResponse{
-					Executions: []*p.InternalVisibilityWorkflowExecutionInfo{
-						{
-							DomainID: DomainID,
-						},
-					},
-				}, fmt.Errorf("error")).Times(1)
-				_, err := visibilityStore.GetClosedWorkflowExecution(context.Background(), test.request)
-				assert.Error(t, err)
+				assert.Equal(t, test.expectedError.Error(), err.Error())
 			} else {
-				mockPinotClient.EXPECT().GetTableName().Return(testTableName).Times(1)
-				mockPinotClient.EXPECT().Search(gomock.Any()).Return(&pnt.SearchResponse{
-					Executions: []*p.InternalVisibilityWorkflowExecutionInfo{
-						{
-							DomainID: DomainID,
-						},
-					},
-				}, nil).Times(1)
-				_, err := visibilityStore.GetClosedWorkflowExecution(context.Background(), test.request)
 				assert.NoError(t, err)
 			}
 		})
@@ -803,18 +889,26 @@ func TestListWorkflowExecutions(t *testing.T) {
 	request := &p.ListWorkflowExecutionsByQueryRequest{}
 
 	tests := map[string]struct {
-		request       *p.ListWorkflowExecutionsByQueryRequest
-		expectedResp  *p.InternalListWorkflowExecutionsResponse
-		expectedError error
+		request                   *p.ListWorkflowExecutionsByQueryRequest
+		expectedResp              *p.InternalListWorkflowExecutionsResponse
+		pinotClientMockAffordance func(mockPinotClient *pnt.MockGenericClient)
+		expectedError             error
 	}{
 		"Case1: error case": {
-			request:       errorRequest,
-			expectedResp:  nil,
-			expectedError: fmt.Errorf("error"),
+			request:      errorRequest,
+			expectedResp: nil,
+			pinotClientMockAffordance: func(mockPinotClient *pnt.MockGenericClient) {
+				mockPinotClient.EXPECT().GetTableName().Return(testTableName).Times(1)
+			},
+			expectedError: fmt.Errorf("next page token: unable to deserialize page token. err: invalid character 'e' looking for beginning of value"),
 		},
 		"Case2: normal case with nil response": {
-			request:       request,
-			expectedResp:  nil,
+			request:      request,
+			expectedResp: nil,
+			pinotClientMockAffordance: func(mockPinotClient *pnt.MockGenericClient) {
+				mockPinotClient.EXPECT().GetTableName().Return(testTableName).Times(1)
+				mockPinotClient.EXPECT().Search(gomock.Any()).Return(nil, nil).Times(1)
+			},
 			expectedError: nil,
 		},
 	}
@@ -830,16 +924,13 @@ func TestListWorkflowExecutions(t *testing.T) {
 			}, mockProducer, log.NewNoop())
 			visibilityStore := mgr.(*pinotVisibilityStore)
 
+			test.pinotClientMockAffordance(mockPinotClient)
+
+			resp, err := visibilityStore.ListWorkflowExecutions(context.Background(), test.request)
+			assert.Equal(t, test.expectedResp, resp)
 			if test.expectedError != nil {
-				mockPinotClient.EXPECT().GetTableName().Return(testTableName).Times(1)
-				resp, err := visibilityStore.ListWorkflowExecutions(context.Background(), test.request)
-				assert.Equal(t, test.expectedResp, resp)
-				assert.Error(t, err)
+				assert.Equal(t, test.expectedError.Error(), err.Error())
 			} else {
-				mockPinotClient.EXPECT().GetTableName().Return(testTableName).Times(1)
-				mockPinotClient.EXPECT().Search(gomock.Any()).Return(nil, nil).Times(1)
-				resp, err := visibilityStore.ListWorkflowExecutions(context.Background(), test.request)
-				assert.Equal(t, test.expectedResp, resp)
 				assert.NoError(t, err)
 			}
 		})
@@ -853,18 +944,26 @@ func TestScanWorkflowExecutions(t *testing.T) {
 	request := &p.ListWorkflowExecutionsByQueryRequest{}
 
 	tests := map[string]struct {
-		request       *p.ListWorkflowExecutionsByQueryRequest
-		expectedResp  *p.InternalListWorkflowExecutionsResponse
-		expectedError error
+		request                   *p.ListWorkflowExecutionsByQueryRequest
+		expectedResp              *p.InternalListWorkflowExecutionsResponse
+		pinotClientMockAffordance func(mockPinotClient *pnt.MockGenericClient)
+		expectedError             error
 	}{
 		"Case1: error case": {
-			request:       errorRequest,
-			expectedResp:  nil,
-			expectedError: fmt.Errorf("error"),
+			request:      errorRequest,
+			expectedResp: nil,
+			pinotClientMockAffordance: func(mockPinotClient *pnt.MockGenericClient) {
+				mockPinotClient.EXPECT().GetTableName().Return(testTableName).Times(1)
+			},
+			expectedError: fmt.Errorf("next page token: unable to deserialize page token. err: invalid character 'e' looking for beginning of value"),
 		},
 		"Case2: normal case with nil response": {
-			request:       request,
-			expectedResp:  nil,
+			request:      request,
+			expectedResp: nil,
+			pinotClientMockAffordance: func(mockPinotClient *pnt.MockGenericClient) {
+				mockPinotClient.EXPECT().GetTableName().Return(testTableName).Times(1)
+				mockPinotClient.EXPECT().Search(gomock.Any()).Return(nil, nil).Times(1)
+			},
 			expectedError: nil,
 		},
 	}
@@ -880,16 +979,13 @@ func TestScanWorkflowExecutions(t *testing.T) {
 			}, mockProducer, log.NewNoop())
 			visibilityStore := mgr.(*pinotVisibilityStore)
 
+			test.pinotClientMockAffordance(mockPinotClient)
+
+			resp, err := visibilityStore.ScanWorkflowExecutions(context.Background(), test.request)
+			assert.Equal(t, test.expectedResp, resp)
 			if test.expectedError != nil {
-				mockPinotClient.EXPECT().GetTableName().Return(testTableName).Times(1)
-				resp, err := visibilityStore.ScanWorkflowExecutions(context.Background(), test.request)
-				assert.Equal(t, test.expectedResp, resp)
-				assert.Error(t, err)
+				assert.Equal(t, test.expectedError.Error(), err.Error())
 			} else {
-				mockPinotClient.EXPECT().GetTableName().Return(testTableName).Times(1)
-				mockPinotClient.EXPECT().Search(gomock.Any()).Return(nil, nil).Times(1)
-				resp, err := visibilityStore.ScanWorkflowExecutions(context.Background(), test.request)
-				assert.Equal(t, test.expectedResp, resp)
 				assert.NoError(t, err)
 			}
 		})
@@ -901,18 +997,27 @@ func TestCountWorkflowExecutions(t *testing.T) {
 	request := &p.CountWorkflowExecutionsRequest{}
 
 	tests := map[string]struct {
-		request       *p.CountWorkflowExecutionsRequest
-		expectedResp  *p.CountWorkflowExecutionsResponse
-		expectedError error
+		request                   *p.CountWorkflowExecutionsRequest
+		expectedResp              *p.CountWorkflowExecutionsResponse
+		pinotClientMockAffordance func(mockPinotClient *pnt.MockGenericClient)
+		expectedError             error
 	}{
 		"Case1: error case": {
-			request:       errorRequest,
-			expectedResp:  nil,
-			expectedError: fmt.Errorf("error"),
+			request:      errorRequest,
+			expectedResp: nil,
+			pinotClientMockAffordance: func(mockPinotClient *pnt.MockGenericClient) {
+				mockPinotClient.EXPECT().GetTableName().Return(testTableName).Times(1)
+				mockPinotClient.EXPECT().CountByQuery(gomock.Any()).Return(int64(0), fmt.Errorf("error")).Times(1)
+			},
+			expectedError: fmt.Errorf("CountClosedWorkflowExecutions failed, error"),
 		},
 		"Case2: normal case with nil response": {
-			request:       request,
-			expectedResp:  &p.CountWorkflowExecutionsResponse{Count: 1},
+			request:      request,
+			expectedResp: &p.CountWorkflowExecutionsResponse{Count: 1},
+			pinotClientMockAffordance: func(mockPinotClient *pnt.MockGenericClient) {
+				mockPinotClient.EXPECT().GetTableName().Return(testTableName).Times(1)
+				mockPinotClient.EXPECT().CountByQuery(gomock.Any()).Return(int64(1), nil).Times(1)
+			},
 			expectedError: nil,
 		},
 	}
@@ -928,17 +1033,13 @@ func TestCountWorkflowExecutions(t *testing.T) {
 			}, mockProducer, log.NewNoop())
 			visibilityStore := mgr.(*pinotVisibilityStore)
 
+			test.pinotClientMockAffordance(mockPinotClient)
+
+			resp, err := visibilityStore.CountWorkflowExecutions(context.Background(), test.request)
+			assert.Equal(t, test.expectedResp, resp)
 			if test.expectedError != nil {
-				mockPinotClient.EXPECT().GetTableName().Return(testTableName).Times(1)
-				mockPinotClient.EXPECT().CountByQuery(gomock.Any()).Return(int64(0), fmt.Errorf("error")).Times(1)
-				resp, err := visibilityStore.CountWorkflowExecutions(context.Background(), test.request)
-				assert.Equal(t, test.expectedResp, resp)
-				assert.Error(t, err)
+				assert.Equal(t, test.expectedError.Error(), err.Error())
 			} else {
-				mockPinotClient.EXPECT().GetTableName().Return(testTableName).Times(1)
-				mockPinotClient.EXPECT().CountByQuery(gomock.Any()).Return(int64(1), nil).Times(1)
-				resp, err := visibilityStore.CountWorkflowExecutions(context.Background(), test.request)
-				assert.Equal(t, test.expectedResp, resp)
 				assert.NoError(t, err)
 			}
 		})