Function TestBatchFlowExecution

pocketflow_test.go:174–202

(t *testing.T)

}

func TestBatchFlowExecution(t *testing.T) {
	batchFlow := pf.NewBatchFlow(simpleLogNode()) // Start node logs based on params

	batchFlow.SetPrepBatch(func(ctx *pf.PfContext, params map[string]any) ([]map[string]any, error) {
		// Generate parameter sets for each batch run
		return []map[string]any{
			{"multiplier": 2},
			{"multiplier": 4},
		}, nil
	})

	batchFlow.SetPostBatch(func(ctx *pf.PfContext, params map[string]any, batchPrepResult []map[string]any) (string, error) {
		// 原代码：ctx["postBatchCalled"] = true
		ctx.SetValue("postBatchCalled", true)
		assert.Len(t, batchPrepResult, 2, "PostBatch should receive the original prep result")
		return "batch_complete", nil
	})

	batchContext := pf.WithParam(context.Background(), nil)
	resultAction, err := batchFlow.Run(batchContext)
	require.NoError(t, err)

	assert.Equal(t, "batch_complete", resultAction)
	assert.True(t, batchContext.Value("postBatchCalled").(bool))

	// Check that the log messages were stored in the shared context by the simpleLogNode's PostFunc
	assert.Equal(t, "SimpleLogNode executed with multiplier: 2", batchContext.Value("last_message_from_batch_2"))
	assert.Equal(t, "SimpleLogNode executed with multiplier: 4", batchContext.Value("last_message_from_batch_4"))
}

// --- Additional Tests ---