diff --git a/submitqueue/orchestrator/controller/batch/batch.go b/submitqueue/orchestrator/controller/batch/batch.go index 88c37f4c..b2742045 100644 --- a/submitqueue/orchestrator/controller/batch/batch.go +++ b/submitqueue/orchestrator/controller/batch/batch.go @@ -106,6 +106,7 @@ func (c *Controller) Process(ctx context.Context, delivery consumer.Delivery) er // Short-circuit if the request has been halted — either it already reached a // terminal state, or the cancel controller has recorded a cancellation intent // (RequestStateCancelling). A halted request must never spawn a new batch. + // If cancellation races with an attempt already initializing below, speculate re-checks the contained request state before starting work. if entity.IsRequestStateHalted(request.State) { metrics.NamedCounter(c.metricsScope, opName, "skipped_halted", 1) c.logger.Infow("skipping batch for halted request", @@ -128,7 +129,7 @@ func (c *Controller) Process(ctx context.Context, delivery consumer.Delivery) er ID: fmt.Sprintf("%s/batch/%d", request.Queue, seq), Queue: request.Queue, Contains: []string{request.ID}, - State: entity.BatchStateCreated, + State: entity.BatchStateCreating, Version: 1, } @@ -167,36 +168,6 @@ func (c *Controller) Process(ctx context.Context, delivery consumer.Delivery) er batch.Dependencies = conflictingIDs - // Update reverse index for each conflicting batch (BatchDependent = - // "batches that depend on me"). One UpdateDependents call per conflict. - for _, depID := range conflictingIDs { - existing, err := c.store.GetBatchDependentStore().Get(ctx, depID) - if err != nil { - metrics.NamedCounter(c.metricsScope, opName, "batch_dependent_store_errors", 1) - return fmt.Errorf("failed to get batch dependent for batchID=%s: %w", depID, err) - } - - dependents := append(existing.Dependents, batch.ID) - - newVersion := existing.Version + 1 - if err := c.store.GetBatchDependentStore().UpdateDependents(ctx, depID, existing.Version, newVersion, dependents); err != nil { - metrics.NamedCounter(c.metricsScope, opName, "batch_dependent_store_errors", 1) - return fmt.Errorf("failed to update batch dependent index for existing batchID=%s and new batchID=%s: %w", depID, batch.ID, err) - } - } - - // Create new reverse index entry for the new batch. It would be empty for now, but will be updated as new batches are created that conflict with this batch. - bd := entity.BatchDependent{ - BatchID: batch.ID, - Dependents: []string{}, - Version: 1, - } - - if err := c.store.GetBatchDependentStore().Create(ctx, bd); err != nil { - metrics.NamedCounter(c.metricsScope, opName, "batch_dependent_store_errors", 1) - return fmt.Errorf("failed to create batch dependent index for new batchID=%s: %w", batch.ID, err) - } - // Claim the request for this batch with a CAS-write that transitions the // request to RequestStateBatched. This CAS is the serialization point // between the batch controller and the cancel controller — without it, the @@ -223,9 +194,8 @@ func (c *Controller) Process(ctx context.Context, delivery consumer.Delivery) er // RequestStateBatched) and cancel.markCancelling(... RequestStateCancelling) // reaches storage first wins; the loser sees storage.ErrVersionMismatch: // - If cancel won: this CAS fails. We ack the message (cancel will drive R - // to its terminal state on its own; no batch is needed). The reverse-index - // entry above becomes a dangling BatchDependent — tolerated per the - // "downstream should handle stale entries" contract on this store. + // to its terminal state on its own; no batch or reverse-index data has + // been written). // - If batch won: cancel.markCancelling will fail with ErrVersionMismatch // on its next attempt, re-fetch R, observe RequestStateBatched, and take // the batch-cancellation branch (which terminates the whole batch). @@ -270,14 +240,19 @@ func (c *Controller) Process(ctx context.Context, delivery consumer.Delivery) er request.Version = newRequestVersion request.State = entity.RequestStateBatched - // Persist batch to storage. - // This is the final operation that concludes the batch creation process. If it fails, BatchDependents will be pointing to a batch id that does not exist. - // We do not reuse batch ids, a retry of this operation will create a new batch with a new ID. The downstream logic that operates on BatchDependent should be able to handle stale entries. + // Persist the batch before creating references to it. A Creating batch is not eligible for dependency analysis or normal processing. if err := c.store.GetBatchStore().Create(ctx, batch); err != nil { metrics.NamedCounter(c.metricsScope, opName, "batch_store_errors", 1) return fmt.Errorf("failed to create batch in batch store: %w", err) } + batch, err = c.initializeBatch(ctx, batch) + if err != nil { + // Retries intentionally mint a new batch ID. Failures may therefore leave unpublished Creating or Created attempts behind. + // These attempts are inert and can be removed by a future background cleanup job if their volume becomes significant. + return err + } + c.logger.Infow("batch created", "batch_id", batch.ID, "request_id", request.ID, @@ -315,6 +290,43 @@ func (c *Controller) Process(ctx context.Context, delivery consumer.Delivery) er return nil // Success - message will be acked } +// initializeBatch creates the reverse-index structure and marks a Creating batch ready for publication. +func (c *Controller) initializeBatch(ctx context.Context, batch entity.Batch) (entity.Batch, error) { + batchDependent := entity.BatchDependent{ + BatchID: batch.ID, + Dependents: []string{}, + Version: 1, + } + if err := c.store.GetBatchDependentStore().Create(ctx, batchDependent); err != nil { + metrics.NamedCounter(c.metricsScope, opName, "batch_dependent_store_errors", 1) + return entity.Batch{}, fmt.Errorf("failed to create batch dependent index for new batchID=%s: %w", batch.ID, err) + } + + for _, dependencyID := range batch.Dependencies { + existing, err := c.store.GetBatchDependentStore().Get(ctx, dependencyID) + if err != nil { + metrics.NamedCounter(c.metricsScope, opName, "batch_dependent_store_errors", 1) + return entity.Batch{}, fmt.Errorf("failed to get batch dependent for batchID=%s: %w", dependencyID, err) + } + + dependents := append(existing.Dependents, batch.ID) + newVersion := existing.Version + 1 + if err := c.store.GetBatchDependentStore().UpdateDependents(ctx, dependencyID, existing.Version, newVersion, dependents); err != nil { + metrics.NamedCounter(c.metricsScope, opName, "batch_dependent_store_errors", 1) + return entity.Batch{}, fmt.Errorf("failed to update batch dependent index for existing batchID=%s and new batchID=%s: %w", dependencyID, batch.ID, err) + } + } + + newVersion := batch.Version + 1 + if err := c.store.GetBatchStore().UpdateState(ctx, batch.ID, batch.Version, newVersion, entity.BatchStateCreated); err != nil { + metrics.NamedCounter(c.metricsScope, opName, "batch_store_errors", 1) + return entity.Batch{}, fmt.Errorf("failed to mark batch %s created: %w", batch.ID, err) + } + batch.Version = newVersion + batch.State = entity.BatchStateCreated + return batch, nil +} + // publish publishes a batch ID to the specified topic key. func (c *Controller) publish(ctx context.Context, key consumer.TopicKey, batchID string, partitionKey string) error { bid := entity.BatchID{ID: batchID} diff --git a/submitqueue/orchestrator/controller/batch/batch_test.go b/submitqueue/orchestrator/controller/batch/batch_test.go index fdfef8e0..32fe238e 100644 --- a/submitqueue/orchestrator/controller/batch/batch_test.go +++ b/submitqueue/orchestrator/controller/batch/batch_test.go @@ -86,6 +86,7 @@ func newTestController(t *testing.T, ctrl *gomock.Controller, cnt *countermock.M mockBatchStore := storagemock.NewMockBatchStore(ctrl) mockBatchStore.EXPECT().GetByQueueAndStates(gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, nil).AnyTimes() mockBatchStore.EXPECT().Create(gomock.Any(), gomock.Any()).Return(nil).AnyTimes() + mockBatchStore.EXPECT().UpdateState(gomock.Any(), gomock.Any(), int32(1), int32(2), entity.BatchStateCreated).Return(nil).AnyTimes() mockReqStore := storagemock.NewMockRequestStore(ctrl) req := testRequest() @@ -168,6 +169,7 @@ func TestController_Process_PublishesBatchedLog(t *testing.T) { mockBatchStore := storagemock.NewMockBatchStore(ctrl) mockBatchStore.EXPECT().GetByQueueAndStates(gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, nil).AnyTimes() mockBatchStore.EXPECT().Create(gomock.Any(), gomock.Any()).Return(nil) + mockBatchStore.EXPECT().UpdateState(gomock.Any(), "test-queue/batch/1", int32(1), int32(2), entity.BatchStateCreated).Return(nil) mockReqStore := storagemock.NewMockRequestStore(ctrl) mockReqStore.EXPECT().Get(gomock.Any(), request.ID).Return(request, nil) @@ -299,6 +301,7 @@ func TestController_Process_WithDependencies(t *testing.T) { mockBatchStore := storagemock.NewMockBatchStore(ctrl) mockBatchStore.EXPECT().GetByQueueAndStates(gomock.Any(), "test-queue", gomock.Any()).Return(activeBatches, nil) mockBatchStore.EXPECT().Create(gomock.Any(), gomock.Any()).Return(nil) + mockBatchStore.EXPECT().UpdateState(gomock.Any(), "test-queue/batch/1", int32(1), int32(2), entity.BatchStateCreated).Return(nil) mockBatchDependentStore := storagemock.NewMockBatchDependentStore(ctrl) // batch/1 has no existing dependents. @@ -351,6 +354,7 @@ func TestController_Process_AnalyzerSelectsSubset(t *testing.T) { mockBatchStore := storagemock.NewMockBatchStore(ctrl) mockBatchStore.EXPECT().GetByQueueAndStates(gomock.Any(), "test-queue", gomock.Any()).Return(activeBatches, nil) mockBatchStore.EXPECT().Create(gomock.Any(), gomock.Any()).Return(nil) + mockBatchStore.EXPECT().UpdateState(gomock.Any(), "test-queue/batch/1", int32(1), int32(2), entity.BatchStateCreated).Return(nil) mockBatchDependentStore := storagemock.NewMockBatchDependentStore(ctrl) // Only batch/2 is selected by the analyzer, so only it gets a reverse-index update. @@ -494,11 +498,6 @@ func TestController_Process_CASLostToCancel(t *testing.T) { mockBatchStore.EXPECT().GetByQueueAndStates(gomock.Any(), "test-queue", gomock.Any()).Return(nil, nil) // Create must NOT be called — gomock fails if it is. - mockBatchDependentStore := storagemock.NewMockBatchDependentStore(ctrl) - // The reverse-index Create still runs because it precedes the CAS; this is - // tolerated per the "downstream handles stale entries" contract. - mockBatchDependentStore.EXPECT().Create(gomock.Any(), gomock.Any()).Return(nil) - mockReqStore := storagemock.NewMockRequestStore(ctrl) mockReqStore.EXPECT().Get(gomock.Any(), request.ID).Return(request, nil) mockReqStore.EXPECT().UpdateState( @@ -507,7 +506,6 @@ func TestController_Process_CASLostToCancel(t *testing.T) { mockStorage := storagemock.NewMockStorage(ctrl) mockStorage.EXPECT().GetBatchStore().Return(mockBatchStore).AnyTimes() - mockStorage.EXPECT().GetBatchDependentStore().Return(mockBatchDependentStore).AnyTimes() mockStorage.EXPECT().GetRequestStore().Return(mockReqStore).AnyTimes() // Publisher with no EXPECTs — must not be called. @@ -547,9 +545,6 @@ func TestController_Process_CASUnexpectedErrorPropagates(t *testing.T) { mockBatchStore.EXPECT().GetByQueueAndStates(gomock.Any(), "test-queue", gomock.Any()).Return(nil, nil) // Create must NOT be called — gomock fails if it is. - mockBatchDependentStore := storagemock.NewMockBatchDependentStore(ctrl) - mockBatchDependentStore.EXPECT().Create(gomock.Any(), gomock.Any()).Return(nil) - casErr := fmt.Errorf("db connection lost") mockReqStore := storagemock.NewMockRequestStore(ctrl) mockReqStore.EXPECT().Get(gomock.Any(), request.ID).Return(request, nil) @@ -559,7 +554,6 @@ func TestController_Process_CASUnexpectedErrorPropagates(t *testing.T) { mockStorage := storagemock.NewMockStorage(ctrl) mockStorage.EXPECT().GetBatchStore().Return(mockBatchStore).AnyTimes() - mockStorage.EXPECT().GetBatchDependentStore().Return(mockBatchDependentStore).AnyTimes() mockStorage.EXPECT().GetRequestStore().Return(mockReqStore).AnyTimes() controller := newTestController(t, ctrl, newSequentialCounter(ctrl), mockStorage, nil, nil) @@ -592,6 +586,7 @@ func TestController_Process_RecoveryAfterPriorCAS(t *testing.T) { mockBatchStore := storagemock.NewMockBatchStore(ctrl) mockBatchStore.EXPECT().GetByQueueAndStates(gomock.Any(), "test-queue", gomock.Any()).Return(nil, nil) mockBatchStore.EXPECT().Create(gomock.Any(), gomock.Any()).Return(nil) + mockBatchStore.EXPECT().UpdateState(gomock.Any(), "test-queue/batch/1", int32(1), int32(2), entity.BatchStateCreated).Return(nil) mockBatchDependentStore := storagemock.NewMockBatchDependentStore(ctrl) mockBatchDependentStore.EXPECT().Create(gomock.Any(), gomock.Any()).Return(nil) @@ -616,3 +611,227 @@ func TestController_Process_RecoveryAfterPriorCAS(t *testing.T) { require.NoError(t, controller.Process(context.Background(), delivery)) } + +func TestController_Process_ReadiesBatchBeforePublishing(t *testing.T) { + ctrl := gomock.NewController(t) + + request := testRequest() + batch := entity.Batch{ + ID: "test-queue/batch/7", + Queue: request.Queue, + Contains: []string{request.ID}, + Dependencies: []string{}, + State: entity.BatchStateCreating, + Version: 1, + } + + cnt := countermock.NewMockCounter(ctrl) + cnt.EXPECT().Next(gomock.Any(), "batch/"+request.Queue).Return(int64(7), nil) + + requestStore := storagemock.NewMockRequestStore(ctrl) + requestStore.EXPECT().Get(gomock.Any(), request.ID).Return(request, nil) + + batchStore := storagemock.NewMockBatchStore(ctrl) + batchStore.EXPECT().GetByQueueAndStates(gomock.Any(), request.Queue, entity.DependencyBatchStates()).Return(nil, nil) + + batchDependentStore := storagemock.NewMockBatchDependentStore(ctrl) + publisher := queuemock.NewMockPublisher(ctrl) + gomock.InOrder( + requestStore.EXPECT().UpdateState(gomock.Any(), request.ID, int32(1), int32(2), entity.RequestStateBatched).Return(nil), + batchStore.EXPECT().Create(gomock.Any(), batch).Return(nil), + batchDependentStore.EXPECT().Create(gomock.Any(), entity.BatchDependent{ + BatchID: batch.ID, + Dependents: []string{}, + Version: 1, + }).Return(nil), + batchStore.EXPECT().UpdateState(gomock.Any(), batch.ID, int32(1), int32(2), entity.BatchStateCreated).Return(nil), + publisher.EXPECT().Publish(gomock.Any(), "log", gomock.Any()).Return(nil), + publisher.EXPECT().Publish(gomock.Any(), "speculate", gomock.Any()).Return(nil), + ) + + store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetRequestStore().Return(requestStore).AnyTimes() + store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() + store.EXPECT().GetBatchDependentStore().Return(batchDependentStore).AnyTimes() + + queue := queuemock.NewMockQueue(ctrl) + queue.EXPECT().Publisher().Return(publisher).AnyTimes() + registry, err := consumer.NewTopicRegistry([]consumer.TopicConfig{ + {Key: topickey.TopicKeySpeculate, Name: "speculate", Queue: queue}, + {Key: topickey.TopicKeyLog, Name: "log", Queue: queue}, + }) + require.NoError(t, err) + + analyzerFactory := conflictmock.NewMockFactory(ctrl) + analyzerFactory.EXPECT().For(conflict.Config{QueueName: request.Queue}).Return(all.New(), nil) + controller := NewController( + zaptest.NewLogger(t).Sugar(), tally.NoopScope, registry, cnt, store, analyzerFactory, + topickey.TopicKeyBatch, "orchestrator-batch", + ) + + msg := entityqueue.NewMessage(request.ID, requestIDPayload(t, request.ID), request.Queue, nil) + delivery := queuemock.NewMockDelivery(ctrl) + delivery.EXPECT().Message().Return(msg).AnyTimes() + delivery.EXPECT().Attempt().Return(1).AnyTimes() + + assert.NoError(t, controller.Process(context.Background(), delivery)) +} + +func TestController_Process_RedeliveryMintsFreshBatchID(t *testing.T) { + ctrl := gomock.NewController(t) + + firstRequest := testRequest() + secondRequest := firstRequest + secondRequest.State = entity.RequestStateBatched + secondRequest.Version = 2 + + cnt := countermock.NewMockCounter(ctrl) + cnt.EXPECT().Next(gomock.Any(), "batch/"+firstRequest.Queue).Return(int64(1), nil) + cnt.EXPECT().Next(gomock.Any(), "batch/"+firstRequest.Queue).Return(int64(2), nil) + + requestStore := storagemock.NewMockRequestStore(ctrl) + requestStore.EXPECT().Get(gomock.Any(), firstRequest.ID).Return(firstRequest, nil) + requestStore.EXPECT().UpdateState(gomock.Any(), firstRequest.ID, int32(1), int32(2), entity.RequestStateBatched).Return(nil) + requestStore.EXPECT().Get(gomock.Any(), firstRequest.ID).Return(secondRequest, nil) + requestStore.EXPECT().UpdateState(gomock.Any(), firstRequest.ID, int32(2), int32(3), entity.RequestStateBatched).Return(nil) + + var createdIDs []string + batchStore := storagemock.NewMockBatchStore(ctrl) + batchStore.EXPECT().GetByQueueAndStates(gomock.Any(), firstRequest.Queue, entity.DependencyBatchStates()).Return(nil, nil).Times(2) + batchStore.EXPECT().Create(gomock.Any(), gomock.Any()).DoAndReturn( + func(_ context.Context, batch entity.Batch) error { + createdIDs = append(createdIDs, batch.ID) + assert.Equal(t, entity.BatchStateCreating, batch.State) + return nil + }, + ).Times(2) + batchStore.EXPECT().UpdateState(gomock.Any(), "test-queue/batch/1", int32(1), int32(2), entity.BatchStateCreated).Return(nil) + batchStore.EXPECT().UpdateState(gomock.Any(), "test-queue/batch/2", int32(1), int32(2), entity.BatchStateCreated).Return(nil) + + batchDependentStore := storagemock.NewMockBatchDependentStore(ctrl) + batchDependentStore.EXPECT().Create(gomock.Any(), gomock.Any()).Return(nil).Times(2) + + store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetRequestStore().Return(requestStore).AnyTimes() + store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() + store.EXPECT().GetBatchDependentStore().Return(batchDependentStore).AnyTimes() + + controller := newTestController(t, ctrl, cnt, store, nil, nil) + msg := entityqueue.NewMessage(firstRequest.ID, requestIDPayload(t, firstRequest.ID), firstRequest.Queue, nil) + delivery := queuemock.NewMockDelivery(ctrl) + delivery.EXPECT().Message().Return(msg).AnyTimes() + delivery.EXPECT().Attempt().Return(2).AnyTimes() + + assert.NoError(t, controller.Process(context.Background(), delivery)) + assert.NoError(t, controller.Process(context.Background(), delivery)) + assert.Equal(t, []string{"test-queue/batch/1", "test-queue/batch/2"}, createdIDs) +} + +func TestController_Process_InitializationFailure(t *testing.T) { + ctrl := gomock.NewController(t) + + request := testRequest() + requestStore := storagemock.NewMockRequestStore(ctrl) + requestStore.EXPECT().Get(gomock.Any(), request.ID).Return(request, nil) + requestStore.EXPECT().UpdateState(gomock.Any(), request.ID, int32(1), int32(2), entity.RequestStateBatched).Return(nil) + + batchStore := storagemock.NewMockBatchStore(ctrl) + batchStore.EXPECT().GetByQueueAndStates(gomock.Any(), request.Queue, entity.DependencyBatchStates()).Return(nil, nil) + batchStore.EXPECT().Create(gomock.Any(), gomock.Any()).Return(nil) + + batchDependentStore := storagemock.NewMockBatchDependentStore(ctrl) + batchDependentStore.EXPECT().Create(gomock.Any(), gomock.Any()).Return(errors.New("storage failed")) + + store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetRequestStore().Return(requestStore).AnyTimes() + store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() + store.EXPECT().GetBatchDependentStore().Return(batchDependentStore).AnyTimes() + + controller := newTestController(t, ctrl, newSequentialCounter(ctrl), store, nil, nil) + msg := entityqueue.NewMessage(request.ID, requestIDPayload(t, request.ID), request.Queue, nil) + delivery := queuemock.NewMockDelivery(ctrl) + delivery.EXPECT().Message().Return(msg).AnyTimes() + delivery.EXPECT().Attempt().Return(1).AnyTimes() + + err := controller.Process(context.Background(), delivery) + assert.ErrorContains(t, err, "failed to create batch dependent index") +} + +func TestInitializeBatchErrors(t *testing.T) { + batch := entity.Batch{ + ID: "test-queue/batch/1", + Dependencies: []string{"test-queue/batch/0"}, + State: entity.BatchStateCreating, + Version: 1, + } + storeErr := errors.New("storage failed") + + tests := []struct { + name string + mockFunc func(*storagemock.MockBatchStore, *storagemock.MockBatchDependentStore) + errMsg string + }{ + { + name: "own reverse index create fails", + mockFunc: func(_ *storagemock.MockBatchStore, dependentStore *storagemock.MockBatchDependentStore) { + dependentStore.EXPECT().Create(gomock.Any(), gomock.Any()).Return(storeErr) + }, + errMsg: "failed to create batch dependent index", + }, + { + name: "dependency get fails", + mockFunc: func(_ *storagemock.MockBatchStore, dependentStore *storagemock.MockBatchDependentStore) { + dependentStore.EXPECT().Create(gomock.Any(), gomock.Any()).Return(nil) + dependentStore.EXPECT().Get(gomock.Any(), "test-queue/batch/0").Return(entity.BatchDependent{}, storeErr) + }, + errMsg: "failed to get batch dependent", + }, + { + name: "dependency update fails", + mockFunc: func(_ *storagemock.MockBatchStore, dependentStore *storagemock.MockBatchDependentStore) { + dependentStore.EXPECT().Create(gomock.Any(), gomock.Any()).Return(nil) + dependentStore.EXPECT().Get(gomock.Any(), "test-queue/batch/0").Return(entity.BatchDependent{ + BatchID: "test-queue/batch/0", + Version: 2, + }, nil) + dependentStore.EXPECT().UpdateDependents( + gomock.Any(), "test-queue/batch/0", int32(2), int32(3), []string{batch.ID}, + ).Return(storeErr) + }, + errMsg: "failed to update batch dependent index", + }, + { + name: "created transition fails", + mockFunc: func(batchStore *storagemock.MockBatchStore, dependentStore *storagemock.MockBatchDependentStore) { + dependentStore.EXPECT().Create(gomock.Any(), gomock.Any()).Return(nil) + dependentStore.EXPECT().Get(gomock.Any(), "test-queue/batch/0").Return(entity.BatchDependent{ + BatchID: "test-queue/batch/0", + Dependents: []string{"test-queue/batch/old"}, + Version: 2, + }, nil) + dependentStore.EXPECT().UpdateDependents( + gomock.Any(), "test-queue/batch/0", int32(2), int32(3), []string{"test-queue/batch/old", batch.ID}, + ).Return(nil) + batchStore.EXPECT().UpdateState(gomock.Any(), batch.ID, int32(1), int32(2), entity.BatchStateCreated).Return(storeErr) + }, + errMsg: "failed to mark batch", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + ctrl := gomock.NewController(t) + batchStore := storagemock.NewMockBatchStore(ctrl) + batchDependentStore := storagemock.NewMockBatchDependentStore(ctrl) + tt.mockFunc(batchStore, batchDependentStore) + + store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() + store.EXPECT().GetBatchDependentStore().Return(batchDependentStore).AnyTimes() + + controller := &Controller{metricsScope: tally.NoopScope, store: store} + _, err := controller.initializeBatch(context.Background(), batch) + assert.ErrorContains(t, err, tt.errMsg) + }) + } +}