Fix OnWorkflow mock to see propagated context headers

internal/internal_workflow_testsuite.go

@@ -2046,6 +2046,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

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())
+}