From 3126433df826071d52ab106ceae33a11b2a8cd8a Mon Sep 17 00:00:00 2001 From: Kriti Behl Date: Mon, 2 Mar 2026 16:29:26 -0500 Subject: [PATCH 1/2] Fix OnWorkflow mock to see propagated context headers --- internal/internal_workflow_testsuite.go | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/internal/internal_workflow_testsuite.go b/internal/internal_workflow_testsuite.go index 17c80b5bc..0cdaf0b7d 100644 --- a/internal/internal_workflow_testsuite.go +++ b/internal/internal_workflow_testsuite.go @@ -2038,6 +2038,11 @@ func (w *workflowExecutorWrapper) Execute(ctx Context, input *commonpb.Payloads) if err != nil { return nil, err } + // Ensure ctxCopy matches real execution: apply header propagation to the context + ctxCopy, err = workflowContextWithHeaderPropagated(ctxCopy, w.env.header, w.env.GetContextPropagators()) + if err != nil { + return nil, err + } go func() { // getWorkflowMockReturn could block if mock is configured to wait. The returned mockRet is what has been configured // for the mock by using MockCallWrapper.Return(). The mockRet could be mock values or mock function. We process From 1a579be356ac000d7280010f7a3f941730fe4a22 Mon Sep 17 00:00:00 2001 From: Kriti Behl Date: Fri, 6 Mar 2026 16:57:01 -0500 Subject: [PATCH 2/2] Add regression test for OnWorkflow mock header propagation --- internal/internal_workflow_testsuite_test.go | 37 ++++++++++++++++++++ 1 file changed, 37 insertions(+) diff --git a/internal/internal_workflow_testsuite_test.go b/internal/internal_workflow_testsuite_test.go index 03e629d28..a8ef3a116 100644 --- a/internal/internal_workflow_testsuite_test.go +++ b/internal/internal_workflow_testsuite_test.go @@ -4915,3 +4915,40 @@ func (s *WorkflowTestSuiteUnitTest) TestChannelWorkerPattern() { s.ElementsMatch([]int{1, 2, 3}, received) }) } +func (s *WorkflowTestSuiteUnitTest) Test_OnWorkflowMockSeesHeaderContext() { + headerSeen := false + + childWorkflowFn := func(ctx Context) error { + return nil + } + + workflowFn := func(ctx Context) error { + cwo := ChildWorkflowOptions{WorkflowRunTimeout: time.Hour} + ctx = WithChildWorkflowOptions(ctx, cwo) + return ExecuteChildWorkflow(ctx, childWorkflowFn).Get(ctx, nil) + } + + env := s.NewTestWorkflowEnvironment() + env.SetHeader(&commonpb.Header{ + Fields: map[string]*commonpb.Payload{ + testHeader: encodeString(s.T(), "test-data"), + }, + }) + env.SetContextPropagators([]ContextPropagator{NewKeysPropagator([]string{testHeader})}) + env.RegisterWorkflow(childWorkflowFn) + + env.OnWorkflow(childWorkflowFn, mock.MatchedBy(func(ctx Context) bool { + val := ctx.Value(contextKey(testHeader)) + if v, ok := val.(string); ok && v == "test-data" { + headerSeen = true + } + return true + })).Return(nil) + + env.ExecuteWorkflow(workflowFn) + + s.True(env.IsWorkflowCompleted()) + s.NoError(env.GetWorkflowError()) + s.True(headerSeen, "OnWorkflow mock should see propagated header in context") + env.AssertExpectations(s.T()) +} \ No newline at end of file