WaitForCancellation activityOption does not work in WorkflowTestSuite

[kittinanasi]

Describe the bug
In WorkflowTestSuite unit tests, when the current workflow is cancelled by calling the CancelWorkflow() on the TestWorkflowEnvironment, if the WaitForCancellation activityOption is ignored when cancelling the workflow.

To Reproduce
Is the issue reproducible?

• Yes

Steps to reproduce the behavior:
The following unit test should succeed:

import (
	"context"
	"go.uber.org/cadence/testsuite"
	"go.uber.org/cadence/workflow"

	"go.uber.org/zap/zaptest"
	"testing"
	"time"

	"github.com/stretchr/testify/mock"
	"github.com/stretchr/testify/suite"
)

type WaitForCancellationTestSuite struct {
	suite.Suite
	testsuite.WorkflowTestSuite
}

func (s *WaitForCancellationTestSuite) SetupSuite() {
}

func (s *WaitForCancellationTestSuite) SetupTest() {
	s.SetLogger(zaptest.NewLogger(s.T()))
}

func TestWaitForCancellationUnitTestSuite(t *testing.T) {
	suite.Run(t, new(WaitForCancellationTestSuite))
}

func (s *WaitForCancellationTestSuite) Test_ActivityWaitForCancellation() {
	env := s.NewTestWorkflowEnvironment()
	env.RegisterWorkflow(testWorkflowHello)
	env.RegisterActivity(testActivityHello)

	c := make(chan struct{})
	activityFinished := false
	env.OnActivity(testActivityHello, mock.Anything, mock.Anything).
		Run(func(args mock.Arguments) {
			c <- struct{}{}
			time.Sleep(time.Second)
			activityFinished = true
		}).
		Return("mock_value", nil).Once()

	go func() {
		<-c
		env.CancelWorkflow()
	}()

	env.ExecuteWorkflow(testWorkflowHello)

	s.True(env.IsWorkflowCompleted())
	s.NoError(env.GetWorkflowError())
	var result string
	env.GetWorkflowResult(&result)
	s.Equal("mock_value", result)

	s.True(activityFinished)
	env.AssertExpectations(s.T())
}

func testWorkflowHello(ctx workflow.Context) (string, error) {
	ao := workflow.ActivityOptions{
		ScheduleToStartTimeout: time.Minute,
		StartToCloseTimeout:    time.Minute,
		HeartbeatTimeout:       20 * time.Second,
		WaitForCancellation:    true,
	}
	ctx = workflow.WithActivityOptions(ctx, ao)

	var result string
	err := workflow.ExecuteActivity(ctx, testActivityHello, "world").Get(ctx, &result)
	if err != nil {
		return "", err
	}
	return result, nil
}

func testActivityHello(_ context.Context, msg string) (string, error) {
	return "hello" + "_" + msg, nil
}

Expected behavior
The above unit test should succeed, because when cancelling the workflow, the workflow execution should wait for the ongoing activity to finish.

Screenshots

Additional context
I think it would be beneficial if the WorkflowTestSuite didn't ignore the WaitForCancellation ActivityOption, so that workflow cancellations could be better tested with unit tests.