diff --git a/tests/testcore/dynamic_config_overrides.go b/tests/testcore/dynamic_config_overrides.go index f08db10a7a..1d3ba7b24b 100644 --- a/tests/testcore/dynamic_config_overrides.go +++ b/tests/testcore/dynamic_config_overrides.go @@ -27,6 +27,7 @@ var ( // unique namespaces with overrides per namespace should be used for tests that require overrides. defaultDynamicConfigOverrides = map[dynamicconfig.Key]any{ dynamicconfig.FrontendRPS.Key(): 3000, + dynamicconfig.FrontendNamespaceReplicationInducingAPIsRPS.Key(): 1000, dynamicconfig.FrontendMaxNamespaceVisibilityRPSPerInstance.Key(): 50, dynamicconfig.FrontendMaxNamespaceVisibilityBurstRatioPerInstance.Key(): 1, dynamicconfig.ReplicationTaskProcessorErrorRetryMaxAttempts.Key(): 1, diff --git a/tests/versioning_3_test.go b/tests/versioning_3_test.go index 4b57300901..66722006e2 100644 --- a/tests/versioning_3_test.go +++ b/tests/versioning_3_test.go @@ -9,11 +9,9 @@ import ( "testing" "time" - "github.com/dgryski/go-farm" "github.com/google/uuid" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" - "github.com/stretchr/testify/suite" commandpb "go.temporal.io/api/command/v1" commonpb "go.temporal.io/api/common/v1" deploymentpb "go.temporal.io/api/deployment/v1" @@ -39,10 +37,12 @@ import ( "go.temporal.io/server/common/dynamicconfig" "go.temporal.io/server/common/primitives/timestamp" "go.temporal.io/server/common/searchattribute/sadefs" + "go.temporal.io/server/common/testing/parallelsuite" "go.temporal.io/server/common/testing/protoutils" "go.temporal.io/server/common/testing/taskpoller" "go.temporal.io/server/common/testing/testhooks" "go.temporal.io/server/common/testing/testvars" + "go.temporal.io/server/common/testing/updateutils" "go.temporal.io/server/common/tqid" "go.temporal.io/server/common/worker_versioning" "go.temporal.io/server/service/matching" @@ -53,8 +53,6 @@ import ( "google.golang.org/protobuf/types/known/timestamppb" ) -var _ = testhooks.MatchingIgnoreRoutingConfigRevisionCheck - type versionStatus int const ( @@ -72,124 +70,135 @@ const ( versionStatusRamping = versionStatus(3) versionStatusDraining = versionStatus(4) versionStatusDrained = versionStatus(5) + + versioning3DeploymentWorkflowVersion = workerdeployment.VersionDataRevisionNumber ) +var _ = testhooks.MatchingIgnoreRoutingConfigRevisionCheck + type Versioning3Suite struct { - testcore.FunctionalTestBase - deploymentWorkflowVersion workerdeployment.DeploymentWorkflowVersion + parallelsuite.Suite[*Versioning3Suite] } func TestVersioning3FunctionalSuite(t *testing.T) { - t.Parallel() - suite.Run(t, &Versioning3Suite{ - deploymentWorkflowVersion: workerdeployment.VersionDataRevisionNumber, - }) + parallelsuite.Run(t, &Versioning3Suite{}) } -func (s *Versioning3Suite) SetupSuite() { - dynamicConfigOverrides := map[dynamicconfig.Key]any{ - dynamicconfig.MatchingDeploymentWorkflowVersion.Key(): int(s.deploymentWorkflowVersion), +func (s *Versioning3Suite) setupEnv(opts ...testcore.TestOption) *testcore.TestEnv { + opts = append([]testcore.TestOption{ + testcore.WithDynamicConfig(dynamicconfig.MatchingDeploymentWorkflowVersion, int(versioning3DeploymentWorkflowVersion)), // Make sure we don't hit the rate limiter in tests - dynamicconfig.FrontendGlobalNamespaceNamespaceReplicationInducingAPIsRPS.Key(): 1000, - dynamicconfig.FrontendMaxNamespaceNamespaceReplicationInducingAPIsBurstRatioPerInstance.Key(): 1, - dynamicconfig.FrontendNamespaceReplicationInducingAPIsRPS.Key(): 1000, + testcore.WithDynamicConfig(dynamicconfig.FrontendGlobalNamespaceNamespaceReplicationInducingAPIsRPS, 1000), + testcore.WithDynamicConfig(dynamicconfig.FrontendMaxNamespaceNamespaceReplicationInducingAPIsBurstRatioPerInstance, 1), - // this is overridden for tests using RunTestWithMatchingBehavior - dynamicconfig.MatchingNumTaskqueueReadPartitions.Key(): 4, - dynamicconfig.MatchingNumTaskqueueWritePartitions.Key(): 4, + // This is overridden for tests using runTestWithMatchingBehavior. + testcore.WithDynamicConfig(dynamicconfig.MatchingNumTaskqueueReadPartitions, 4), + testcore.WithDynamicConfig(dynamicconfig.MatchingNumTaskqueueWritePartitions, 4), // Overriding the number of deployments that can be registered in a single namespace. Done only for this test suite // since it creates a large number of unique deployments in the test suite's namespace. - dynamicconfig.MatchingMaxDeployments.Key(): 1000, + testcore.WithDynamicConfig(dynamicconfig.MatchingMaxDeployments, 1000), + }, opts...) + + return testcore.NewEnv(s.T(), opts...) +} + +func (s *Versioning3Suite) runTestWithMatchingBehavior(testFn func(*testcore.TestEnv, *Versioning3Suite), opts ...testcore.TestOption) { + for _, behavior := range testcore.AllMatchingBehaviors() { + behavior := behavior + s.Run(behavior.Name(), func(s *Versioning3Suite) { + envOpts := append([]testcore.TestOption{}, opts...) + envOpts = append(envOpts, behavior.Options()...) + env := s.setupEnv(envOpts...) + behavior.InjectHooks(env) + testFn(env, s) + }) } - s.FunctionalTestBase.SetupSuiteWithCluster(testcore.WithDynamicConfigOverrides(dynamicConfigOverrides)) +} + +func (s *Versioning3Suite) updateUtils() updateutils.UpdateUtils { + return updateutils.New(s.T()) } func (s *Versioning3Suite) TestPinnedTask_NoProperPoller() { - s.RunTestWithMatchingBehavior( - func() { - tv := testvars.New(s) + s.runTestWithMatchingBehavior(func(env *testcore.TestEnv, s *Versioning3Suite) { + tv := env.Tv() - tv2 := tv.WithBuildIDNumber(2) - go s.idlePollWorkflow(context.Background(), tv2, true, ver3MinPollTime, "second deployment should not receive pinned task") + tv2 := tv.WithBuildIDNumber(2) + go s.idlePollWorkflow(env, env.Context(), tv2, true, ver3MinPollTime, "second deployment should not receive pinned task") - // Start a versioned poller for the first version so that it registers the version in the task queue. - pollerCtx, cancelPoller := context.WithCancel(context.Background()) - go s.idlePollWorkflow(pollerCtx, tv, true, ver3MinPollTime, "first deployment should not receive any task. It is just creating a version in the task queue.") + // Start a versioned poller for the first version so that it registers the version in the task queue. + pollerCtx, cancelPoller := context.WithCancel(env.Context()) + go s.idlePollWorkflow(env, pollerCtx, tv, true, ver3MinPollTime, "first deployment should not receive any task. It is just creating a version in the task queue.") - // Wait for the version to be present in the task queue - s.validatePinnedVersionExistsInTaskQueue(tv) + // Wait for the version to be present in the task queue + s.validatePinnedVersionExistsInTaskQueue(env, tv) - // Cancel the poller after condition is met - cancelPoller() + // Cancel the poller after condition is met + cancelPoller() - s.startWorkflow(tv, tv.VersioningOverridePinned()) - s.idlePollWorkflow(context.Background(), tv, false, ver3MinPollTime, "unversioned worker should not receive pinned task") + s.startWorkflow(env, tv, tv.VersioningOverridePinned()) + s.idlePollWorkflow(env, env.Context(), tv, false, ver3MinPollTime, "unversioned worker should not receive pinned task") - // Sleeping to let the pollers arrive to server before ending the test. - time.Sleep(200 * time.Millisecond) //nolint:forbidigo - }) + // Sleeping to let the pollers arrive to server before ending the test. + time.Sleep(200 * time.Millisecond) //nolint:forbidigo + }) } func (s *Versioning3Suite) TestUnpinnedTask_NonCurrentDeployment() { - s.RunTestWithMatchingBehavior( - func() { - tv := testvars.New(s) - go s.idlePollWorkflow(context.Background(), tv, true, ver3MinPollTime, "non-current versioned poller should not receive unpinned task") + s.runTestWithMatchingBehavior(func(env *testcore.TestEnv, s *Versioning3Suite) { + tv := env.Tv() + go s.idlePollWorkflow(env, env.Context(), tv, true, ver3MinPollTime, "non-current versioned poller should not receive unpinned task") - s.startWorkflow(tv, nil) + s.startWorkflow(env, tv, nil) - // Sleeping to let the pollers arrive to server before ending the test. - time.Sleep(200 * time.Millisecond) //nolint:forbidigo - }) + // Sleeping to let the pollers arrive to server before ending the test. + time.Sleep(200 * time.Millisecond) //nolint:forbidigo + }) } func (s *Versioning3Suite) TestUnpinnedTask_OldDeployment() { - s.RunTestWithMatchingBehavior( - func() { - tv := testvars.New(s) - tvOldDeployment := tv.WithBuildIDNumber(1) - tvNewDeployment := tv.WithBuildIDNumber(2) - - // previous current deployment - s.updateTaskQueueDeploymentDataWithRoutingConfig(tvOldDeployment, &deploymentpb.RoutingConfig{ - CurrentDeploymentVersion: worker_versioning.ExternalWorkerDeploymentVersionFromStringV31(tvOldDeployment.DeploymentVersionString()), - CurrentVersionChangedTime: timestamp.TimePtr(time.Now()), - RevisionNumber: 1, - }, map[string]*deploymentspb.WorkerDeploymentVersionData{tvOldDeployment.DeploymentVersion().GetBuildId(): { - Status: enumspb.WORKER_DEPLOYMENT_VERSION_STATUS_CURRENT, - }}, []string{}, tqTypeWf) - - // current deployment - s.updateTaskQueueDeploymentDataWithRoutingConfig(tvNewDeployment, &deploymentpb.RoutingConfig{ - CurrentDeploymentVersion: worker_versioning.ExternalWorkerDeploymentVersionFromStringV31(tvNewDeployment.DeploymentVersionString()), - CurrentVersionChangedTime: timestamp.TimePtr(time.Now()), - RevisionNumber: 2, - }, map[string]*deploymentspb.WorkerDeploymentVersionData{tvNewDeployment.DeploymentVersion().GetBuildId(): { - Status: enumspb.WORKER_DEPLOYMENT_VERSION_STATUS_CURRENT, - }}, []string{}, tqTypeWf) + s.runTestWithMatchingBehavior(func(env *testcore.TestEnv, s *Versioning3Suite) { + tv := env.Tv() + tvOldDeployment := tv.WithBuildIDNumber(1) + tvNewDeployment := tv.WithBuildIDNumber(2) + + // previous current deployment + s.updateTaskQueueDeploymentDataWithRoutingConfig(env, tvOldDeployment, &deploymentpb.RoutingConfig{ + CurrentDeploymentVersion: worker_versioning.ExternalWorkerDeploymentVersionFromStringV31(tvOldDeployment.DeploymentVersionString()), + CurrentVersionChangedTime: timestamp.TimePtr(time.Now()), + RevisionNumber: 1, + }, map[string]*deploymentspb.WorkerDeploymentVersionData{tvOldDeployment.DeploymentVersion().GetBuildId(): { + Status: enumspb.WORKER_DEPLOYMENT_VERSION_STATUS_CURRENT, + }}, []string{}, tqTypeWf) - s.startWorkflow(tv, nil) + // current deployment + s.updateTaskQueueDeploymentDataWithRoutingConfig(env, tvNewDeployment, &deploymentpb.RoutingConfig{ + CurrentDeploymentVersion: worker_versioning.ExternalWorkerDeploymentVersionFromStringV31(tvNewDeployment.DeploymentVersionString()), + CurrentVersionChangedTime: timestamp.TimePtr(time.Now()), + RevisionNumber: 2, + }, map[string]*deploymentspb.WorkerDeploymentVersionData{tvNewDeployment.DeploymentVersion().GetBuildId(): { + Status: enumspb.WORKER_DEPLOYMENT_VERSION_STATUS_CURRENT, + }}, []string{}, tqTypeWf) - s.idlePollWorkflow( - context.Background(), - tvOldDeployment, - true, - ver3MinPollTime, - "old deployment should not receive unpinned task", - ) - // Sleeping to let the pollers arrive to server before ending the test. - time.Sleep(200 * time.Millisecond) //nolint:forbidigo - }, - ) + s.startWorkflow(env, tv, nil) + + s.idlePollWorkflow(env, + env.Context(), + tvOldDeployment, + true, + ver3MinPollTime, + "old deployment should not receive unpinned task", + ) + // Sleeping to let the pollers arrive to server before ending the test. + time.Sleep(200 * time.Millisecond) //nolint:forbidigo + }) } func (s *Versioning3Suite) TestSessionActivityResourceSpecificTaskQueueNotRegisteredInVersion() { - tv := testvars.New(s) - - ctx, cancel := context.WithTimeout(context.Background(), time.Minute) - defer cancel() + env := s.setupEnv() + tv := env.Tv() const ( wfName = "session-wf" @@ -221,7 +230,7 @@ func (s *Versioning3Suite) TestSessionActivityResourceSpecificTaskQueueNotRegist return activity.GetInfo(ctx).TaskQueue, nil } - w := worker.New(s.SdkClient(), tv.TaskQueue().GetName(), worker.Options{ + w := worker.New(env.SdkClient(), tv.TaskQueue().GetName(), worker.Options{ DeploymentOptions: worker.DeploymentOptions{ Version: tv.SDKDeploymentVersion(), UseVersioning: true, @@ -235,26 +244,26 @@ func (s *Versioning3Suite) TestSessionActivityResourceSpecificTaskQueueNotRegist defer w.Stop() // Ensure the version is current and propagated before starting the workflow. - s.setCurrentDeployment(tv) - s.waitForDeploymentDataPropagation(tv, versionStatusCurrent, false, tqTypeWf, tqTypeAct) + s.setCurrentDeployment(env, tv) + s.waitForDeploymentDataPropagation(env, tv, versionStatusCurrent, false, tqTypeWf, tqTypeAct) - run, err := s.SdkClient().ExecuteWorkflow(ctx, sdkclient.StartWorkflowOptions{ + run, err := env.SdkClient().ExecuteWorkflow(env.Context(), sdkclient.StartWorkflowOptions{ TaskQueue: tv.TaskQueue().GetName(), }, wfName) s.NoError(err) var sessionTaskQueue string - s.NoError(run.Get(ctx, &sessionTaskQueue)) + s.NoError(run.Get(env.Context(), &sessionTaskQueue)) s.NotEmpty(sessionTaskQueue) // Sanity: for sessions this should be a resource-specific activity task queue, not the base TQ. s.NotEqual(tv.TaskQueue().GetName(), sessionTaskQueue) // The session resource-specific task queue must NOT be registered in the version - resp, err := s.FrontendClient().DescribeWorkerDeploymentVersion(ctx, &workflowservice.DescribeWorkerDeploymentVersionRequest{ - Namespace: s.Namespace().String(), + resp, err := env.FrontendClient().DescribeWorkerDeploymentVersion(env.Context(), &workflowservice.DescribeWorkerDeploymentVersionRequest{ + Namespace: env.Namespace().String(), Version: tv.DeploymentVersionString(), }) - s.Require().NoError(err) + s.NoError(err) totalActTQ := 0 for _, tq := range resp.GetVersionTaskQueues() { @@ -270,156 +279,140 @@ func (s *Versioning3Suite) TestSessionActivityResourceSpecificTaskQueueNotRegist } func (s *Versioning3Suite) TestWorkflowWithPinnedOverride_Sticky() { - s.RunTestWithMatchingBehavior( - func() { - s.testWorkflowWithPinnedOverride(true) - }, - ) + s.runTestWithMatchingBehavior(func(env *testcore.TestEnv, s *Versioning3Suite) { + s.testWorkflowWithPinnedOverride(env, true) + }) } func (s *Versioning3Suite) TestWorkflowWithPinnedOverride_NoSticky() { - s.RunTestWithMatchingBehavior( - func() { - s.testWorkflowWithPinnedOverride(false) - }, - ) + s.runTestWithMatchingBehavior(func(env *testcore.TestEnv, s *Versioning3Suite) { + s.testWorkflowWithPinnedOverride(env, false) + }) } -func (s *Versioning3Suite) testWorkflowWithPinnedOverride(sticky bool) { - ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second) - defer cancel() - - tv := testvars.New(s) +func (s *Versioning3Suite) testWorkflowWithPinnedOverride(env *testcore.TestEnv, sticky bool) { + tv := env.Tv() if sticky { - s.warmUpSticky(tv) + s.warmUpSticky(env, tv) } wftCompleted := make(chan struct{}) - s.pollWftAndHandle(tv, false, wftCompleted, + s.pollWftAndHandle(env, tv, false, wftCompleted, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) - s.verifyVersioningSAs(tv, vbPinned, enumspb.WORKFLOW_EXECUTION_STATUS_RUNNING) + s.verifyVersioningSAs(env, tv, vbPinned, enumspb.WORKFLOW_EXECUTION_STATUS_RUNNING) return respondWftWithActivities(tv, tv, sticky, vbUnpinned, "5"), nil }) actCompleted := make(chan struct{}) - s.pollActivityAndHandle(tv, actCompleted, + s.pollActivityAndHandle(env, tv, actCompleted, func(task *workflowservice.PollActivityTaskQueueResponse) (*workflowservice.RespondActivityTaskCompletedRequest, error) { s.NotNil(task) return respondActivity(), nil }) // Wait for the version to be present in the task queue. Version existence is required before it can be set as an override. - s.validatePinnedVersionExistsInTaskQueue(tv) + s.validatePinnedVersionExistsInTaskQueue(env, tv) - runID := s.startWorkflow(tv, tv.VersioningOverridePinned()) + runID := s.startWorkflow(env, tv, tv.VersioningOverridePinned()) - s.WaitForChannel(ctx, wftCompleted) - s.verifyWorkflowVersioning(s.Assertions, tv, vbUnpinned, tv.Deployment(), tv.VersioningOverridePinned(), nil) - s.verifyVersioningSAs(tv, vbPinned, enumspb.WORKFLOW_EXECUTION_STATUS_RUNNING, tv) + env.WaitForChannel(env.Context(), wftCompleted) + s.verifyWorkflowVersioning(env, s.Assertions, tv, vbUnpinned, tv.Deployment(), tv.VersioningOverridePinned(), nil) + s.verifyVersioningSAs(env, tv, vbPinned, enumspb.WORKFLOW_EXECUTION_STATUS_RUNNING, tv) if sticky { - s.verifyWorkflowStickyQueue(tv.WithRunID(runID)) + s.verifyWorkflowStickyQueue(env, tv.WithRunID(runID)) } - s.WaitForChannel(ctx, actCompleted) - s.verifyWorkflowVersioning(s.Assertions, tv, vbUnpinned, tv.Deployment(), tv.VersioningOverridePinned(), nil) + env.WaitForChannel(env.Context(), actCompleted) + s.verifyWorkflowVersioning(env, s.Assertions, tv, vbUnpinned, tv.Deployment(), tv.VersioningOverridePinned(), nil) - s.pollWftAndHandle(tv, sticky, nil, + s.pollWftAndHandle(env, tv, sticky, nil, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) return respondCompleteWorkflow(tv, vbUnpinned), nil }) - s.verifyWorkflowVersioning(s.Assertions, tv, vbUnpinned, tv.Deployment(), tv.VersioningOverridePinned(), nil) + s.verifyWorkflowVersioning(env, s.Assertions, tv, vbUnpinned, tv.Deployment(), tv.VersioningOverridePinned(), nil) } func (s *Versioning3Suite) TestQueryWithPinnedOverride_NoSticky() { - s.RunTestWithMatchingBehavior( - func() { - s.testQueryWithPinnedOverride(false) - }, - ) + s.runTestWithMatchingBehavior(func(env *testcore.TestEnv, s *Versioning3Suite) { + s.testQueryWithPinnedOverride(env, false) + }) } func (s *Versioning3Suite) TestQueryWithPinnedOverride_Sticky() { - s.RunTestWithMatchingBehavior( - func() { - s.testQueryWithPinnedOverride(true) - }, - ) + s.runTestWithMatchingBehavior(func(env *testcore.TestEnv, s *Versioning3Suite) { + s.testQueryWithPinnedOverride(env, true) + }) } func (s *Versioning3Suite) TestPinnedQuery_DrainedVersion_PollersAbsent() { - s.RunTestWithMatchingBehavior( - func() { - s.testPinnedQuery_DrainedVersion(false, false) - }, + s.runTestWithMatchingBehavior(func(env *testcore.TestEnv, s *Versioning3Suite) { + s.testPinnedQueryDrainedVersion(env, false, false) + }, + testcore.WithDynamicConfig(dynamicconfig.VersionDrainageStatusRefreshInterval, 1*time.Second), + testcore.WithDynamicConfig(dynamicconfig.VersionDrainageStatusVisibilityGracePeriod, 1*time.Second), + testcore.WithDynamicConfig(dynamicconfig.PollerHistoryTTL, 500*time.Millisecond), ) } func (s *Versioning3Suite) TestPinnedQuery_DrainedVersion_PollersPresent() { - s.RunTestWithMatchingBehavior( - func() { - s.testPinnedQuery_DrainedVersion(true, false) - }, + s.runTestWithMatchingBehavior(func(env *testcore.TestEnv, s *Versioning3Suite) { + s.testPinnedQueryDrainedVersion(env, true, false) + }, + testcore.WithDynamicConfig(dynamicconfig.VersionDrainageStatusRefreshInterval, 1*time.Second), + testcore.WithDynamicConfig(dynamicconfig.VersionDrainageStatusVisibilityGracePeriod, 1*time.Second), ) } func (s *Versioning3Suite) TestPinnedQuery_RollbackDrainedVersion() { - s.RunTestWithMatchingBehavior( - func() { - s.testPinnedQuery_DrainedVersion(true, true) - }, + s.runTestWithMatchingBehavior(func(env *testcore.TestEnv, s *Versioning3Suite) { + s.testPinnedQueryDrainedVersion(env, true, true) + }, + testcore.WithDynamicConfig(dynamicconfig.VersionDrainageStatusRefreshInterval, 1*time.Second), + testcore.WithDynamicConfig(dynamicconfig.VersionDrainageStatusVisibilityGracePeriod, 1*time.Second), ) } -func (s *Versioning3Suite) testPinnedQuery_DrainedVersion(pollersPresent bool, rollback bool) { - ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) - defer cancel() - - s.OverrideDynamicConfig(dynamicconfig.VersionDrainageStatusRefreshInterval, 1*time.Second) - s.OverrideDynamicConfig(dynamicconfig.VersionDrainageStatusVisibilityGracePeriod, 1*time.Second) - if !pollersPresent { - s.OverrideDynamicConfig(dynamicconfig.PollerHistoryTTL, 500*time.Millisecond) - } - - tv := testvars.New(s) +func (s *Versioning3Suite) testPinnedQueryDrainedVersion(env *testcore.TestEnv, pollersPresent bool, rollback bool) { + tv := env.Tv() // create version v1 and make it current idlePollerDone := make(chan struct{}) go func() { - s.idlePollWorkflow(context.Background(), tv, true, ver3MinPollTime, "should not have gotten any tasks since there are none") + s.idlePollWorkflow(env, env.Context(), tv, true, ver3MinPollTime, "should not have gotten any tasks since there are none") close(idlePollerDone) }() - s.setCurrentDeployment(tv) - s.WaitForChannel(ctx, idlePollerDone) + s.setCurrentDeployment(env, tv) + env.WaitForChannel(env.Context(), idlePollerDone) wftCompleted := make(chan struct{}) - s.pollWftAndHandle(tv, false, wftCompleted, + s.pollWftAndHandle(env, tv, false, wftCompleted, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) return respondCompleteWorkflow(tv, vbPinned), nil }) - s.startWorkflow(tv, tv.VersioningOverridePinned()) - s.WaitForChannel(ctx, wftCompleted) - s.verifyWorkflowVersioning(s.Assertions, tv, vbPinned, tv.Deployment(), tv.VersioningOverridePinned(), nil) + s.startWorkflow(env, tv, tv.VersioningOverridePinned()) + env.WaitForChannel(env.Context(), wftCompleted) + s.verifyWorkflowVersioning(env, s.Assertions, tv, vbPinned, tv.Deployment(), tv.VersioningOverridePinned(), nil) // create version v2 and make it current which shall make v1 go from current -> draining/drained idlePollerDone = make(chan struct{}) tv2 := tv.WithBuildIDNumber(2) go func() { - s.idlePollWorkflow(context.Background(), tv2, true, ver3MinPollTime, "should not have gotten any tasks since there are none") + s.idlePollWorkflow(env, env.Context(), tv2, true, ver3MinPollTime, "should not have gotten any tasks since there are none") close(idlePollerDone) }() - s.setCurrentDeployment(tv2) - s.WaitForChannel(ctx, idlePollerDone) + s.setCurrentDeployment(env, tv2) + env.WaitForChannel(env.Context(), idlePollerDone) // wait for v1 to become drained s.EventuallyWithT(func(t *assert.CollectT) { a := require.New(t) - resp, err := s.FrontendClient().DescribeWorkerDeploymentVersion(ctx, &workflowservice.DescribeWorkerDeploymentVersionRequest{ - Namespace: s.Namespace().String(), + resp, err := env.FrontendClient().DescribeWorkerDeploymentVersion(env.Context(), &workflowservice.DescribeWorkerDeploymentVersionRequest{ + Namespace: env.Namespace().String(), Version: tv.DeploymentVersionString(), }) a.NoError(err) @@ -430,27 +423,27 @@ func (s *Versioning3Suite) testPinnedQuery_DrainedVersion(pollersPresent bool, r // simulate the pollers going away, which should make the query fail as now the version is drained + has no pollers polling it time.Sleep(1 * time.Second) //nolint:forbidigo versionStr := tv.Deployment().GetBuildId() - if s.deploymentWorkflowVersion < workerdeployment.AsyncSetCurrentAndRamping { + if versioning3DeploymentWorkflowVersion < workerdeployment.AsyncSetCurrentAndRamping { versionStr = worker_versioning.ExternalWorkerDeploymentVersionToString(worker_versioning.ExternalWorkerDeploymentVersionFromDeployment(tv.Deployment())) } - _, err := s.queryWorkflow(tv) + _, err := s.queryWorkflow(env, tv) s.Error(err) s.ErrorContains(err, fmt.Sprintf(matching.ErrBlackholedQuery, versionStr, versionStr)) } else { // since the version still has pollers, the query should succeed - s.pollAndQueryWorkflow(tv, false) + s.pollAndQueryWorkflow(env, tv, false) } if rollback { // ramp traffic back to v1 to remove it from drained state - s.setRampingDeployment(tv, 50, false) + s.setRampingDeployment(env, tv, 50, false) // wait for v1 to become ramping s.EventuallyWithT(func(t *assert.CollectT) { a := require.New(t) - resp, err := s.FrontendClient().DescribeWorkerDeploymentVersion(ctx, &workflowservice.DescribeWorkerDeploymentVersionRequest{ - Namespace: s.Namespace().String(), + resp, err := env.FrontendClient().DescribeWorkerDeploymentVersion(env.Context(), &workflowservice.DescribeWorkerDeploymentVersionRequest{ + Namespace: env.Namespace().String(), Version: tv.DeploymentVersionString(), }) a.NoError(err) @@ -458,26 +451,23 @@ func (s *Versioning3Suite) testPinnedQuery_DrainedVersion(pollersPresent bool, r }, time.Second*10, time.Millisecond*1000) // the ramping status is propagated to the task queues - s.waitForDeploymentDataPropagation(tv, versionStatusRamping, false, tqTypeWf) + s.waitForDeploymentDataPropagation(env, tv, versionStatusRamping, false, tqTypeWf) // the version has pollers and is ramping making the query succeed - s.pollAndQueryWorkflow(tv, false) + s.pollAndQueryWorkflow(env, tv, false) } } -func (s *Versioning3Suite) testQueryWithPinnedOverride(sticky bool) { - ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second) - defer cancel() - - tv := testvars.New(s) +func (s *Versioning3Suite) testQueryWithPinnedOverride(env *testcore.TestEnv, sticky bool) { + tv := env.Tv() if sticky { - s.warmUpSticky(tv) + s.warmUpSticky(env, tv) } wftCompleted := make(chan struct{}) - s.pollWftAndHandle(tv, false, wftCompleted, + s.pollWftAndHandle(env, tv, false, wftCompleted, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) return respondEmptyWft(tv, sticky, vbUnpinned), nil @@ -486,8 +476,8 @@ func (s *Versioning3Suite) testQueryWithPinnedOverride(sticky bool) { // Wait for the version to be present in the task queue. Version existence is required before it can be set as an override. s.EventuallyWithT(func(t *assert.CollectT) { a := require.New(t) - resp, err := s.GetTestCluster().MatchingClient().CheckTaskQueueVersionMembership(ctx, &matchingservice.CheckTaskQueueVersionMembershipRequest{ - NamespaceId: s.NamespaceID().String(), + resp, err := env.GetTestCluster().MatchingClient().CheckTaskQueueVersionMembership(env.Context(), &matchingservice.CheckTaskQueueVersionMembershipRequest{ + NamespaceId: env.NamespaceID().String(), TaskQueue: tv.TaskQueue().GetName(), TaskQueueType: tqTypeWf, Version: worker_versioning.DeploymentVersionFromDeployment(tv.Deployment()), @@ -496,103 +486,95 @@ func (s *Versioning3Suite) testQueryWithPinnedOverride(sticky bool) { a.True(resp.GetIsMember()) }, 10*time.Second, 100*time.Millisecond) - runID := s.startWorkflow(tv, tv.VersioningOverridePinned()) + runID := s.startWorkflow(env, tv, tv.VersioningOverridePinned()) - s.WaitForChannel(ctx, wftCompleted) - s.verifyWorkflowVersioning(s.Assertions, tv, vbUnpinned, tv.Deployment(), tv.VersioningOverridePinned(), nil) + env.WaitForChannel(env.Context(), wftCompleted) + s.verifyWorkflowVersioning(env, s.Assertions, tv, vbUnpinned, tv.Deployment(), tv.VersioningOverridePinned(), nil) if sticky { - s.verifyWorkflowStickyQueue(tv.WithRunID(runID)) + s.verifyWorkflowStickyQueue(env, tv.WithRunID(runID)) } - s.pollAndQueryWorkflow(tv, sticky) + s.pollAndQueryWorkflow(env, tv, sticky) } func (s *Versioning3Suite) TestUnpinnedQuery_NoSticky() { - s.RunTestWithMatchingBehavior( - func() { - s.testUnpinnedQuery(false) - }, - ) + s.runTestWithMatchingBehavior(func(env *testcore.TestEnv, s *Versioning3Suite) { + s.testUnpinnedQuery(env, false) + }) } func (s *Versioning3Suite) TestUnpinnedQuery_Sticky() { - s.RunTestWithMatchingBehavior( - func() { - s.testUnpinnedQuery(true) - }, - ) + s.runTestWithMatchingBehavior(func(env *testcore.TestEnv, s *Versioning3Suite) { + s.testUnpinnedQuery(env, true) + }) } -func (s *Versioning3Suite) testUnpinnedQuery(sticky bool) { - ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second) - defer cancel() - - tv := testvars.New(s) +func (s *Versioning3Suite) testUnpinnedQuery(env *testcore.TestEnv, sticky bool) { + tv := env.Tv() tv2 := tv.WithBuildIDNumber(2) if sticky { - s.warmUpSticky(tv) + s.warmUpSticky(env, tv) } wftCompleted := make(chan struct{}) - s.pollWftAndHandle(tv, false, wftCompleted, + s.pollWftAndHandle(env, tv, false, wftCompleted, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) - s.verifyWorkflowVersioning(s.Assertions, tv, vbUnspecified, nil, nil, tv.DeploymentVersionTransition()) + s.verifyWorkflowVersioning(env, s.Assertions, tv, vbUnspecified, nil, nil, tv.DeploymentVersionTransition()) return respondEmptyWft(tv, sticky, vbUnpinned), nil }) - s.setCurrentDeployment(tv) - s.waitForDeploymentDataPropagation(tv, versionStatusCurrent, false, tqTypeWf) + s.setCurrentDeployment(env, tv) + s.waitForDeploymentDataPropagation(env, tv, versionStatusCurrent, false, tqTypeWf) - runID := s.startWorkflow(tv, nil) + runID := s.startWorkflow(env, tv, nil) - s.WaitForChannel(ctx, wftCompleted) - s.verifyWorkflowVersioning(s.Assertions, tv, vbUnpinned, tv.Deployment(), nil, nil) + env.WaitForChannel(env.Context(), wftCompleted) + s.verifyWorkflowVersioning(env, s.Assertions, tv, vbUnpinned, tv.Deployment(), nil, nil) if sticky { - s.verifyWorkflowStickyQueue(tv.WithRunID(runID)) + s.verifyWorkflowStickyQueue(env, tv.WithRunID(runID)) } pollerDone := make(chan struct{}) go func() { - s.idlePollWorkflow(context.Background(), tv2, true, 5*time.Second, "new deployment should not receive query") + s.idlePollWorkflow(env, env.Context(), tv2, true, 5*time.Second, "new deployment should not receive query") close(pollerDone) }() - s.pollAndQueryWorkflow(tv, sticky) - s.WaitForChannel(ctx, pollerDone) // wait for the idle poller to complete to not interfere with the next poller + s.pollAndQueryWorkflow(env, tv, sticky) + env.WaitForChannel(env.Context(), pollerDone) // wait for the idle poller to complete to not interfere with the next poller - s.setCurrentDeployment(tv2) - s.waitForDeploymentDataPropagation(tv2, versionStatusCurrent, false, tqTypeWf) + s.setCurrentDeployment(env, tv2) + s.waitForDeploymentDataPropagation(env, tv2, versionStatusCurrent, false, tqTypeWf) - go s.idlePollWorkflow(context.Background(), tv, true, ver3MinPollTime, "old deployment should not receive query") + go s.idlePollWorkflow(env, env.Context(), tv, true, ver3MinPollTime, "old deployment should not receive query") // Since the current deployment has changed, task will move to the normal queue (thus, sticky=false) - s.pollAndQueryWorkflow(tv2, false) + s.pollAndQueryWorkflow(env, tv2, false) } func (s *Versioning3Suite) pollAndQueryWorkflow( + env *testcore.TestEnv, tv *testvars.TestVars, sticky bool, ) { queryResultCh := make(chan any) - s.pollWftAndHandleQueries(tv, sticky, queryResultCh, + s.pollWftAndHandleQueries(env, tv, sticky, queryResultCh, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondQueryTaskCompletedRequest, error) { return &workflowservice.RespondQueryTaskCompletedRequest{}, nil }) - _, err := s.queryWorkflow(tv) + _, err := s.queryWorkflow(env, tv) s.NoError(err) <-queryResultCh } func (s *Versioning3Suite) TestPinnedWorkflowWithLateActivityPoller() { - s.RunTestWithMatchingBehavior( - func() { - s.testPinnedWorkflowWithLateActivityPoller() - }, - ) + s.runTestWithMatchingBehavior(func(env *testcore.TestEnv, s *Versioning3Suite) { + s.testPinnedWorkflowWithLateActivityPoller(env) + }) } -func (s *Versioning3Suite) testPinnedWorkflowWithLateActivityPoller() { +func (s *Versioning3Suite) testPinnedWorkflowWithLateActivityPoller(env *testcore.TestEnv) { // Here, we test that designating activities as independent is revisited if the missing activity // pollers arrive to server while the so-far-independent activity is backlogged. // Summary: a wf starts with a pinned override. The first wft schedules an activity before @@ -600,121 +582,110 @@ func (s *Versioning3Suite) testPinnedWorkflowWithLateActivityPoller() { // to the default queue. Then, the activity poller on the pinned deployment arrives, the task // should be now sent to that poller although no current deployment is set on the TQs. - ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second) - defer cancel() - - tv := testvars.New(s) + tv := env.Tv() wftCompleted := make(chan struct{}) - s.pollWftAndHandle(tv, false, wftCompleted, + s.pollWftAndHandle(env, tv, false, wftCompleted, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) return respondWftWithActivities(tv, tv, false, vbUnpinned, "5"), nil }) - s.waitForDeploymentDataPropagation(tv, versionStatusInactive, false, tqTypeWf) + s.waitForDeploymentDataPropagation(env, tv, versionStatusInactive, false, tqTypeWf) override := tv.VersioningOverridePinned() - s.startWorkflow(tv, override) + s.startWorkflow(env, tv, override) - s.WaitForChannel(ctx, wftCompleted) - s.verifyWorkflowVersioning(s.Assertions, tv, vbUnpinned, tv.Deployment(), override, nil) + env.WaitForChannel(env.Context(), wftCompleted) + s.verifyWorkflowVersioning(env, s.Assertions, tv, vbUnpinned, tv.Deployment(), override, nil) // Wait long enough to make sure the activity is backlogged. - s.validateBacklogCount(tv, tqTypeAct, 1) + s.validateBacklogCount(env, tv, tqTypeAct, 1) // When the first activity poller arrives from this deployment, it registers the TQ in the // deployment and that will trigger reevaluation of backlog queue. - s.pollActivityAndHandle(tv, nil, + s.pollActivityAndHandle(env, tv, nil, func(task *workflowservice.PollActivityTaskQueueResponse) (*workflowservice.RespondActivityTaskCompletedRequest, error) { s.NotNil(task) return respondActivity(), nil }) - s.verifyWorkflowVersioning(s.Assertions, tv, vbUnpinned, tv.Deployment(), override, nil) - s.validateBacklogCount(tv, tqTypeAct, 0) + s.verifyWorkflowVersioning(env, s.Assertions, tv, vbUnpinned, tv.Deployment(), override, nil) + s.validateBacklogCount(env, tv, tqTypeAct, 0) - s.pollWftAndHandle(tv, false, nil, + s.pollWftAndHandle(env, tv, false, nil, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) return respondCompleteWorkflow(tv, vbUnpinned), nil }) - s.verifyWorkflowVersioning(s.Assertions, tv, vbUnpinned, tv.Deployment(), override, nil) + s.verifyWorkflowVersioning(env, s.Assertions, tv, vbUnpinned, tv.Deployment(), override, nil) } func (s *Versioning3Suite) TestUnpinnedWorkflow_Sticky() { - s.RunTestWithMatchingBehavior( - func() { - s.testUnpinnedWorkflow(true) - }, - ) + s.runTestWithMatchingBehavior(func(env *testcore.TestEnv, s *Versioning3Suite) { + s.testUnpinnedWorkflow(env, true) + }) } func (s *Versioning3Suite) TestUnpinnedWorkflow_NoSticky() { - s.RunTestWithMatchingBehavior( - func() { - s.testUnpinnedWorkflow(false) - }, - ) + s.runTestWithMatchingBehavior(func(env *testcore.TestEnv, s *Versioning3Suite) { + s.testUnpinnedWorkflow(env, false) + }) } -func (s *Versioning3Suite) testUnpinnedWorkflow(sticky bool) { - ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second) - defer cancel() - - tv := testvars.New(s) +func (s *Versioning3Suite) testUnpinnedWorkflow(env *testcore.TestEnv, sticky bool) { + tv := env.Tv() if sticky { - s.warmUpSticky(tv) + s.warmUpSticky(env, tv) } wftCompleted := make(chan struct{}) - s.pollWftAndHandle(tv, false, wftCompleted, + s.pollWftAndHandle(env, tv, false, wftCompleted, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) - s.verifyWorkflowVersioning(s.Assertions, tv, vbUnspecified, nil, nil, tv.DeploymentVersionTransition()) + s.verifyWorkflowVersioning(env, s.Assertions, tv, vbUnspecified, nil, nil, tv.DeploymentVersionTransition()) return respondWftWithActivities(tv, tv, sticky, vbUnpinned, "5"), nil }) actCompleted := make(chan struct{}) - s.pollActivityAndHandle(tv, actCompleted, + s.pollActivityAndHandle(env, tv, actCompleted, func(task *workflowservice.PollActivityTaskQueueResponse) (*workflowservice.RespondActivityTaskCompletedRequest, error) { s.NotNil(task) return respondActivity(), nil }) - s.setCurrentDeployment(tv) + s.setCurrentDeployment(env, tv) - runID := s.startWorkflow(tv, nil) + runID := s.startWorkflow(env, tv, nil) - s.WaitForChannel(ctx, wftCompleted) - s.verifyWorkflowVersioning(s.Assertions, tv, vbUnpinned, tv.Deployment(), nil, nil) - s.verifyVersioningSAs(tv, vbUnpinned, enumspb.WORKFLOW_EXECUTION_STATUS_RUNNING, tv) + env.WaitForChannel(env.Context(), wftCompleted) + s.verifyWorkflowVersioning(env, s.Assertions, tv, vbUnpinned, tv.Deployment(), nil, nil) + s.verifyVersioningSAs(env, tv, vbUnpinned, enumspb.WORKFLOW_EXECUTION_STATUS_RUNNING, tv) if sticky { - s.verifyWorkflowStickyQueue(tv.WithRunID(runID)) + s.verifyWorkflowStickyQueue(env, tv.WithRunID(runID)) } - s.WaitForChannel(ctx, actCompleted) - s.verifyWorkflowVersioning(s.Assertions, tv, vbUnpinned, tv.Deployment(), nil, nil) + env.WaitForChannel(env.Context(), actCompleted) + s.verifyWorkflowVersioning(env, s.Assertions, tv, vbUnpinned, tv.Deployment(), nil, nil) - s.pollWftAndHandle(tv, sticky, nil, + s.pollWftAndHandle(env, tv, sticky, nil, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) return respondCompleteWorkflow(tv, vbUnpinned), nil }) - s.verifyWorkflowVersioning(s.Assertions, tv, vbUnpinned, tv.Deployment(), nil, nil) + s.verifyWorkflowVersioning(env, s.Assertions, tv, vbUnpinned, tv.Deployment(), nil, nil) } func (s *Versioning3Suite) TestSearchByUsedVersion() { - ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second) - defer cancel() + env := s.setupEnv() - tv := testvars.New(s) + tv := env.Tv() wftCompleted := make(chan struct{}) - s.pollWftAndHandle(tv, false, wftCompleted, + s.pollWftAndHandle(env, tv, false, wftCompleted, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) // Although pinned to the version by override, the query should not return a workflow because the workflow task is not completed yet. - wfs, err := s.FrontendClient().ListWorkflowExecutions(ctx, &workflowservice.ListWorkflowExecutionsRequest{ - Namespace: s.Namespace().String(), + wfs, err := env.FrontendClient().ListWorkflowExecutions(env.Context(), &workflowservice.ListWorkflowExecutionsRequest{ + Namespace: env.Namespace().String(), Query: fmt.Sprintf("TemporalUsedWorkerDeploymentVersions = '%s'", tv.DeploymentVersionStringV32()), }) s.NoError(err) @@ -722,16 +693,16 @@ func (s *Versioning3Suite) TestSearchByUsedVersion() { return respondCompleteWorkflow(tv, vbPinned), nil }) - s.waitForDeploymentDataPropagation(tv, versionStatusInactive, false, tqTypeWf) - s.startWorkflow(tv, tv.VersioningOverridePinned()) + s.waitForDeploymentDataPropagation(env, tv, versionStatusInactive, false, tqTypeWf) + s.startWorkflow(env, tv, tv.VersioningOverridePinned()) <-wftCompleted - s.verifyVersioningSAs(tv, vbPinned, enumspb.WORKFLOW_EXECUTION_STATUS_COMPLETED, tv) + s.verifyVersioningSAs(env, tv, vbPinned, enumspb.WORKFLOW_EXECUTION_STATUS_COMPLETED, tv) s.EventuallyWithT(func(t *assert.CollectT) { a := assert.New(t) - wfs, err := s.FrontendClient().ListWorkflowExecutions(ctx, &workflowservice.ListWorkflowExecutionsRequest{ - Namespace: s.Namespace().String(), + wfs, err := env.FrontendClient().ListWorkflowExecutions(env.Context(), &workflowservice.ListWorkflowExecutionsRequest{ + Namespace: env.Namespace().String(), Query: fmt.Sprintf("TemporalUsedWorkerDeploymentVersions = '%s'", tv.DeploymentVersionStringV32()), }) a.NoError(err) @@ -743,29 +714,28 @@ func (s *Versioning3Suite) TestSearchByUsedVersion() { // drains the initial workflow task from the execution, and ensures the task is correctly // routed to the appropriate build. func (s *Versioning3Suite) drainWorkflowTaskAfterSetCurrentWithOverride( + env *testcore.TestEnv, tv *testvars.TestVars, override *workflowpb.VersioningOverride, ) (*commonpb.WorkflowExecution, string) { wftCompleted := make(chan struct{}) - s.pollWftAndHandle(tv, false, wftCompleted, + s.pollWftAndHandle(env, tv, false, wftCompleted, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) if override != nil { - s.verifyWorkflowVersioning(s.Assertions, tv, vbUnspecified, nil, override, nil) + s.verifyWorkflowVersioning(env, s.Assertions, tv, vbUnspecified, nil, override, nil) } else { - s.verifyWorkflowVersioning(s.Assertions, tv, vbUnspecified, nil, override, tv.DeploymentVersionTransition()) + s.verifyWorkflowVersioning(env, s.Assertions, tv, vbUnspecified, nil, override, tv.DeploymentVersionTransition()) } return respondEmptyWft(tv, false, vbUnpinned), nil }) - s.waitForDeploymentDataPropagation(tv, versionStatusInactive, false, tqTypeWf) - s.setCurrentDeployment(tv) + s.waitForDeploymentDataPropagation(env, tv, versionStatusInactive, false, tqTypeWf) + s.setCurrentDeployment(env, tv) - runID := s.startWorkflow(tv, override) + runID := s.startWorkflow(env, tv, override) execution := tv.WithRunID(runID).WorkflowExecution() - ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second) - defer cancel() - s.WaitForChannel(ctx, wftCompleted) + env.WaitForChannel(env.Context(), wftCompleted) return execution, runID } @@ -774,43 +744,43 @@ func (s *Versioning3Suite) drainWorkflowTaskAfterSetCurrentWithOverride( // drains the initial workflow task from the execution, and ensures the task is correctly // routed to the appropriate build. func (s *Versioning3Suite) drainWorkflowTaskAfterSetCurrent( + env *testcore.TestEnv, tv *testvars.TestVars, ) (*commonpb.WorkflowExecution, string) { wftCompleted := make(chan struct{}) - s.pollWftAndHandle(tv, false, wftCompleted, + s.pollWftAndHandle(env, tv, false, wftCompleted, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) - s.verifyWorkflowVersioning(s.Assertions, tv, vbUnspecified, nil, nil, tv.DeploymentVersionTransition()) + s.verifyWorkflowVersioning(env, s.Assertions, tv, vbUnspecified, nil, nil, tv.DeploymentVersionTransition()) return respondEmptyWft(tv, false, vbUnpinned), nil }) - s.waitForDeploymentDataPropagation(tv, versionStatusInactive, false, tqTypeWf) - s.setCurrentDeployment(tv) + s.waitForDeploymentDataPropagation(env, tv, versionStatusInactive, false, tqTypeWf) + s.setCurrentDeployment(env, tv) - runID := s.startWorkflow(tv, nil) + runID := s.startWorkflow(env, tv, nil) execution := tv.WithRunID(runID).WorkflowExecution() - ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second) - defer cancel() - s.WaitForChannel(ctx, wftCompleted) + env.WaitForChannel(env.Context(), wftCompleted) return execution, runID } func (s *Versioning3Suite) TestUnpinnedWorkflow_SuccessfulUpdate_TransitionsToNewDeployment() { - tv1 := testvars.New(s).WithBuildIDNumber(1) + env := s.setupEnv() + tv1 := env.Tv().WithBuildIDNumber(1) - execution, _ := s.drainWorkflowTaskAfterSetCurrent(tv1) + execution, _ := s.drainWorkflowTaskAfterSetCurrent(env, tv1) // Register the new version and set it to current tv2 := tv1.WithBuildIDNumber(2) - s.pollUntilRegistered(context.Background(), tv2) - s.setCurrentDeployment(tv2) + s.pollUntilRegistered(env, tv2) + s.setCurrentDeployment(env, tv2) // Send update - updateResultCh := sendUpdateNoError(s, tv2) + updateResultCh := sendUpdateNoError(env, tv2) // Process update in workflow - s.pollWftAndHandle(tv2, false, nil, + s.pollWftAndHandle(env, tv2, false, nil, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) @@ -824,7 +794,7 @@ func (s *Versioning3Suite) TestUnpinnedWorkflow_SuccessfulUpdate_TransitionsToNe 6 WorkflowTaskStarted `, task.History) - events := s.GetHistory(s.Namespace().String(), execution) + events := env.GetHistory(env.Namespace().String(), execution) s.EqualHistoryEvents(` 1 WorkflowExecutionStarted 2 WorkflowTaskScheduled @@ -836,11 +806,11 @@ func (s *Versioning3Suite) TestUnpinnedWorkflow_SuccessfulUpdate_TransitionsToNe // VersioningInfo should not have changed before the update has been processed by the poller. // Deployment version transition should also be nil since this is a speculative task. - s.verifyWorkflowVersioning(s.Assertions, tv1, vbUnpinned, tv1.Deployment(), nil, nil) + s.verifyWorkflowVersioning(env, s.Assertions, tv1, vbUnpinned, tv1.Deployment(), nil, nil) return &workflowservice.RespondWorkflowTaskCompletedRequest{ - Commands: s.UpdateAcceptCompleteCommands(tv2), - Messages: s.UpdateAcceptCompleteMessages(tv2, task.Messages[0]), + Commands: s.updateUtils().UpdateAcceptCompleteCommands(tv2), + Messages: s.updateUtils().UpdateAcceptCompleteMessages(tv2, task.Messages[0]), VersioningBehavior: vbUnpinned, DeploymentOptions: &deploymentpb.WorkerDeploymentOptions{ BuildId: tv2.BuildID(), @@ -854,7 +824,7 @@ func (s *Versioning3Suite) TestUnpinnedWorkflow_SuccessfulUpdate_TransitionsToNe s.Equal("success-result-of-"+tv2.UpdateID(), testcore.DecodeString(s.T(), updateResult.GetOutcome().GetSuccess())) // Verify that events from the speculative task are written to the history since the update was accepted - events := s.GetHistory(s.Namespace().String(), execution) + events := env.GetHistory(env.Namespace().String(), execution) s.EqualHistoryEvents(` 1 WorkflowExecutionStarted 2 WorkflowTaskScheduled @@ -862,14 +832,14 @@ func (s *Versioning3Suite) TestUnpinnedWorkflow_SuccessfulUpdate_TransitionsToNe 4 WorkflowTaskCompleted 5 WorkflowTaskScheduled // Was speculative WT... 6 WorkflowTaskStarted -7 WorkflowTaskCompleted // ...and events were written to the history when WT completes. +7 WorkflowTaskCompleted // ...and events were written to the history when WT completes. 8 WorkflowExecutionUpdateAccepted {"AcceptedRequestSequencingEventId": 5} // WTScheduled event which delivered update to the worker. 9 WorkflowExecutionUpdateCompleted {"AcceptedEventId": 8} `, events) // Verify that the versioning info is updated correctly. - describeCall, err := s.FrontendClient().DescribeWorkflowExecution(context.Background(), &workflowservice.DescribeWorkflowExecutionRequest{ - Namespace: s.Namespace().String(), + describeCall, err := env.FrontendClient().DescribeWorkflowExecution(env.Context(), &workflowservice.DescribeWorkflowExecutionRequest{ + Namespace: env.Namespace().String(), Execution: execution, }) s.NoError(err) @@ -878,25 +848,26 @@ func (s *Versioning3Suite) TestUnpinnedWorkflow_SuccessfulUpdate_TransitionsToNe // Since the poller accepted the update, the Worker Deployment Version that completed the last workflow task // of this workflow execution should have changed to the new version. However, the version transition should // still be nil. - s.verifyWorkflowVersioning(s.Assertions, tv2, vbUnpinned, tv2.Deployment(), nil, nil) + s.verifyWorkflowVersioning(env, s.Assertions, tv2, vbUnpinned, tv2.Deployment(), nil, nil) } func (s *Versioning3Suite) TestUnpinnedWorkflow_FailedUpdate_DoesNotTransitionToNewDeployment() { - tv1 := testvars.New(s).WithBuildIDNumber(1) + env := s.setupEnv() + tv1 := env.Tv().WithBuildIDNumber(1) - execution, _ := s.drainWorkflowTaskAfterSetCurrent(tv1) + execution, _ := s.drainWorkflowTaskAfterSetCurrent(env, tv1) // Register the new version and set it to current tv2 := tv1.WithBuildIDNumber(2) - s.pollUntilRegistered(context.Background(), tv2) - s.setCurrentDeployment(tv2) + s.pollUntilRegistered(env, tv2) + s.setCurrentDeployment(env, tv2) // Send update - updateResultCh := sendUpdateNoError(s, tv2) + updateResultCh := sendUpdateNoError(env, tv2) // Process update in workflow - s.pollWftAndHandle(tv2, false, nil, + s.pollWftAndHandle(env, tv2, false, nil, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) @@ -910,7 +881,7 @@ func (s *Versioning3Suite) TestUnpinnedWorkflow_FailedUpdate_DoesNotTransitionTo 6 WorkflowTaskStarted `, task.History) - events := s.GetHistory(s.Namespace().String(), execution) + events := env.GetHistory(env.Namespace().String(), execution) s.EqualHistoryEvents(` 1 WorkflowExecutionStarted 2 WorkflowTaskScheduled @@ -922,7 +893,7 @@ func (s *Versioning3Suite) TestUnpinnedWorkflow_FailedUpdate_DoesNotTransitionTo // VersioningInfo should not have changed before the update has been processed by the poller. // Deployment version transition should also be nil since this is a speculative task. - s.verifyWorkflowVersioning(s.Assertions, tv1, vbUnpinned, tv1.Deployment(), nil, nil) + s.verifyWorkflowVersioning(env, s.Assertions, tv1, vbUnpinned, tv1.Deployment(), nil, nil) updRequestMsg := task.Messages[0] updRequest := protoutils.UnmarshalAny[*updatepb.Request](s.T(), updRequestMsg.GetBody()) @@ -932,7 +903,7 @@ func (s *Versioning3Suite) TestUnpinnedWorkflow_FailedUpdate_DoesNotTransitionTo s.EqualValues(5, updRequestMsg.GetEventId()) return &workflowservice.RespondWorkflowTaskCompletedRequest{ - Messages: s.UpdateRejectMessages(tv2, updRequestMsg), + Messages: s.updateUtils().UpdateRejectMessages(tv2, updRequestMsg), VersioningBehavior: vbUnpinned, DeploymentOptions: &deploymentpb.WorkerDeploymentOptions{ BuildId: tv2.BuildID(), @@ -946,7 +917,7 @@ func (s *Versioning3Suite) TestUnpinnedWorkflow_FailedUpdate_DoesNotTransitionTo s.Equal("rejection-of-"+tv2.UpdateID(), updateResult.GetOutcome().GetFailure().GetMessage()) // Verify events from the speculative task are *not* written to the workflow history since the update was rejected - events := s.GetHistory(s.Namespace().String(), execution) + events := env.GetHistory(env.Namespace().String(), execution) s.EqualHistoryEvents(` 1 WorkflowExecutionStarted 2 WorkflowTaskScheduled @@ -956,23 +927,19 @@ func (s *Versioning3Suite) TestUnpinnedWorkflow_FailedUpdate_DoesNotTransitionTo // Since the poller rejected the update, the Worker Deployment Version that completed the last workflow task // of this workflow execution should not have changed. - s.verifyWorkflowVersioning(s.Assertions, tv1, vbUnpinned, tv1.Deployment(), nil, nil) + s.verifyWorkflowVersioning(env, s.Assertions, tv1, vbUnpinned, tv1.Deployment(), nil, nil) } func (s *Versioning3Suite) TestUnpinnedWorkflowWithRamp_ToVersioned() { - s.RunTestWithMatchingBehavior( - func() { - s.testUnpinnedWorkflowWithRamp(false) - }, - ) + s.runTestWithMatchingBehavior(func(env *testcore.TestEnv, s *Versioning3Suite) { + s.testUnpinnedWorkflowWithRamp(env, false) + }) } func (s *Versioning3Suite) TestUnpinnedWorkflowWithRamp_ToUnversioned() { - s.RunTestWithMatchingBehavior( - func() { - s.testUnpinnedWorkflowWithRamp(true) - }, - ) + s.runTestWithMatchingBehavior(func(env *testcore.TestEnv, s *Versioning3Suite) { + s.testUnpinnedWorkflowWithRamp(env, true) + }) } func (s *Versioning3Suite) TestWorkflowRetry_Pinned_ExpectInherit_RetryOfChild() { @@ -992,7 +959,9 @@ func (s *Versioning3Suite) TestWorkflowRetry_Unpinned_ExpectAutoUpgradeToCurrent } func (s *Versioning3Suite) testWorkflowRetry(behavior workflow.VersioningBehavior, expectInherit, retryOfChild, retryOfCaN bool) { - tv1 := testvars.New(s).WithBuildIDNumber(1) + env := s.setupEnv() + + tv1 := env.Tv().WithBuildIDNumber(1) tv2 := tv1.WithBuildIDNumber(2) childWorkflowID := tv1.WorkflowID() + "-child" @@ -1031,10 +1000,7 @@ func (s *Versioning3Suite) testWorkflowRetry(behavior workflow.VersioningBehavio return "v2", nil } - ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute) - defer cancel() - - w1 := worker.New(s.SdkClient(), tv1.TaskQueue().GetName(), worker.Options{ + w1 := worker.New(env.SdkClient(), tv1.TaskQueue().GetName(), worker.Options{ DeploymentOptions: worker.DeploymentOptions{ Version: tv1.SDKDeploymentVersion(), UseVersioning: true, @@ -1047,7 +1013,7 @@ func (s *Versioning3Suite) testWorkflowRetry(behavior workflow.VersioningBehavio s.NoError(w1.Start()) defer w1.Stop() - w2 := worker.New(s.SdkClient(), tv1.TaskQueue().GetName(), worker.Options{ + w2 := worker.New(env.SdkClient(), tv1.TaskQueue().GetName(), worker.Options{ DeploymentOptions: worker.DeploymentOptions{ Version: tv2.SDKDeploymentVersion(), UseVersioning: true, @@ -1061,15 +1027,15 @@ func (s *Versioning3Suite) testWorkflowRetry(behavior workflow.VersioningBehavio defer w2.Stop() // Set v1 to current and propagate to all task queue partitions - s.setCurrentDeployment(tv1) - s.waitForDeploymentDataPropagation(tv1, versionStatusCurrent, false, tqTypeWf, tqTypeAct) + s.setCurrentDeployment(env, tv1) + s.waitForDeploymentDataPropagation(env, tv1, versionStatusCurrent, false, tqTypeWf, tqTypeAct) wf0 := "wf" if retryOfChild { wf0 = "parent-wf" } - run0, err := s.SdkClient().ExecuteWorkflow( - ctx, + run0, err := env.SdkClient().ExecuteWorkflow( + env.Context(), sdkclient.StartWorkflowOptions{ TaskQueue: tv1.TaskQueue().GetName(), RetryPolicy: &temporal.RetryPolicy{ @@ -1087,7 +1053,7 @@ func (s *Versioning3Suite) testWorkflowRetry(behavior workflow.VersioningBehavio if retryOfCaN { // wait for first run to continue-as-new s.EventuallyWithT(func(t *assert.CollectT) { - desc, err := s.SdkClient().DescribeWorkflow(ctx, wfIDOfRetryingWF, run0.GetRunID()) + desc, err := env.SdkClient().DescribeWorkflow(env.Context(), wfIDOfRetryingWF, run0.GetRunID()) if !assert.NoError(t, err) { return } @@ -1097,7 +1063,7 @@ func (s *Versioning3Suite) testWorkflowRetry(behavior workflow.VersioningBehavio // wait for workflow to progress on v1 (activity completed and waiting for signal) s.EventuallyWithT(func(t *assert.CollectT) { - desc, err := s.SdkClient().DescribeWorkflowExecution(ctx, wfIDOfRetryingWF, "") + desc, err := env.SdkClient().DescribeWorkflowExecution(env.Context(), wfIDOfRetryingWF, "") if !assert.NoError(t, err) { return } @@ -1110,7 +1076,7 @@ func (s *Versioning3Suite) testWorkflowRetry(behavior workflow.VersioningBehavio wfIDOfRetryingWF = childWorkflowID // Wait for child workflow to be created s.EventuallyWithT(func(t *assert.CollectT) { - desc, err := s.SdkClient().DescribeWorkflow(ctx, wfIDOfRetryingWF, "") + desc, err := env.SdkClient().DescribeWorkflow(env.Context(), wfIDOfRetryingWF, "") if !assert.NoError(t, err) { return } @@ -1119,7 +1085,7 @@ func (s *Versioning3Suite) testWorkflowRetry(behavior workflow.VersioningBehavio } else if retryOfCaN { // get the next run in the continue-as-new chain s.EventuallyWithT(func(t *assert.CollectT) { - continuedAsNewRunResp, err := s.SdkClient().DescribeWorkflow(ctx, wfIDOfRetryingWF, "") + continuedAsNewRunResp, err := env.SdkClient().DescribeWorkflow(env.Context(), wfIDOfRetryingWF, "") if !assert.NoError(t, err) { return } @@ -1132,15 +1098,15 @@ func (s *Versioning3Suite) testWorkflowRetry(behavior workflow.VersioningBehavio } // Set v2 to current and propagate to all task queue partitions - s.setCurrentDeployment(tv2) - s.waitForDeploymentDataPropagation(tv2, versionStatusCurrent, false, tqTypeWf, tqTypeAct) + s.setCurrentDeployment(env, tv2) + s.waitForDeploymentDataPropagation(env, tv2, versionStatusCurrent, false, tqTypeWf, tqTypeAct) // signal workflow to continue (it will fail and then retry on v2 if it doesn't inherit) - s.NoError(s.SdkClient().SignalWorkflow(ctx, wfIDOfRetryingWF, runIDBeforeRetry, "currentVersionChanged", nil)) + s.NoError(env.SdkClient().SignalWorkflow(env.Context(), wfIDOfRetryingWF, runIDBeforeRetry, "currentVersionChanged", nil)) // wait for run that will retry to fail s.EventuallyWithT(func(t *assert.CollectT) { - desc, err := s.SdkClient().DescribeWorkflow(ctx, wfIDOfRetryingWF, runIDBeforeRetry) + desc, err := env.SdkClient().DescribeWorkflow(env.Context(), wfIDOfRetryingWF, runIDBeforeRetry) if !assert.NoError(t, err) { return } @@ -1150,7 +1116,7 @@ func (s *Versioning3Suite) testWorkflowRetry(behavior workflow.VersioningBehavio // get the execution info of the next run in the retry chain, wait for next run to start var secondRunID string s.EventuallyWithT(func(t *assert.CollectT) { - secondRunResp, err := s.SdkClient().DescribeWorkflow(ctx, wfIDOfRetryingWF, "") + secondRunResp, err := env.SdkClient().DescribeWorkflow(env.Context(), wfIDOfRetryingWF, "") if !assert.NoError(t, err) { return } @@ -1161,7 +1127,7 @@ func (s *Versioning3Suite) testWorkflowRetry(behavior workflow.VersioningBehavio // confirm that the second run eventually gets auto-upgrade behavior and runs on version 2 (no inherit) s.EventuallyWithT(func(t *assert.CollectT) { - secondRunResp, err := s.SdkClient().DescribeWorkflowExecution(ctx, wfIDOfRetryingWF, secondRunID) + secondRunResp, err := env.SdkClient().DescribeWorkflowExecution(env.Context(), wfIDOfRetryingWF, secondRunID) if !assert.NoError(t, err) { return } @@ -1186,11 +1152,11 @@ func (s *Versioning3Suite) testWorkflowRetry(behavior workflow.VersioningBehavio }, 5*time.Second, 100*time.Millisecond) } -func (s *Versioning3Suite) testUnpinnedWorkflowWithRamp(toUnversioned bool) { +func (s *Versioning3Suite) testUnpinnedWorkflowWithRamp(env *testcore.TestEnv, toUnversioned bool) { // This test sets a 50% ramp and runs 50 wfs and ensures both versions got some wf and // activity tasks. - tv1 := testvars.New(s).WithBuildIDNumber(1) + tv1 := env.Tv().WithBuildIDNumber(1) tv2 := tv1.WithBuildIDNumber(2) wf := func(ctx workflow.Context, version string) (string, error) { @@ -1219,10 +1185,7 @@ func (s *Versioning3Suite) testUnpinnedWorkflowWithRamp(toUnversioned bool) { return "v2", nil } - ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute) - defer cancel() - - w1 := worker.New(s.SdkClient(), tv1.TaskQueue().GetName(), worker.Options{ + w1 := worker.New(env.SdkClient(), tv1.TaskQueue().GetName(), worker.Options{ DeploymentOptions: worker.DeploymentOptions{ Version: tv1.SDKDeploymentVersion(), UseVersioning: true, @@ -1239,12 +1202,12 @@ func (s *Versioning3Suite) testUnpinnedWorkflowWithRamp(toUnversioned bool) { // Make sure both TQs are registered in v1 which will be the current version. This is to make sure // we don't get to ramping while one of the TQs has not yet got v1 as it's current version. - // (note that s.setCurrentDeployment(tv1) can pass even with one TQ added to the version) - s.waitForDeploymentDataPropagation(tv1, versionStatusInactive, false, tqTypeWf, tqTypeAct) - s.setCurrentDeployment(tv1) + // (note that s.setCurrentDeployment(env, tv1) can pass even with one TQ added to the version) + s.waitForDeploymentDataPropagation(env, tv1, versionStatusInactive, false, tqTypeWf, tqTypeAct) + s.setCurrentDeployment(env, tv1) // wait until all task queue partitions know that tv1 is current - s.waitForDeploymentDataPropagation(tv1, versionStatusCurrent, false, tqTypeWf, tqTypeAct) + s.waitForDeploymentDataPropagation(env, tv1, versionStatusCurrent, false, tqTypeWf, tqTypeAct) deploymentOpts := worker.DeploymentOptions{ Version: tv2.SDKDeploymentVersion(), @@ -1253,7 +1216,7 @@ func (s *Versioning3Suite) testUnpinnedWorkflowWithRamp(toUnversioned bool) { deploymentOpts.UseVersioning = true deploymentOpts.DefaultVersioningBehavior = workflow.VersioningBehaviorAutoUpgrade } - w2 := worker.New(s.SdkClient(), tv2.TaskQueue().GetName(), worker.Options{ + w2 := worker.New(env.SdkClient(), tv2.TaskQueue().GetName(), worker.Options{ DeploymentOptions: deploymentOpts, MaxConcurrentWorkflowTaskPollers: numPollers, }) @@ -1262,21 +1225,21 @@ func (s *Versioning3Suite) testUnpinnedWorkflowWithRamp(toUnversioned bool) { s.NoError(w2.Start()) defer w2.Stop() - s.setRampingDeployment(tv2, 50, toUnversioned) + s.setRampingDeployment(env, tv2, 50, toUnversioned) // wait until all task queue partitions know that tv2 is ramping - s.waitForDeploymentDataPropagation(tv2, versionStatusRamping, toUnversioned, tqTypeWf, tqTypeAct) + s.waitForDeploymentDataPropagation(env, tv2, versionStatusRamping, toUnversioned, tqTypeWf, tqTypeAct) numTests := 50 counter := make(map[string]int) runs := make([]sdkclient.WorkflowRun, numTests) for i := range numTests { - run, err := s.SdkClient().ExecuteWorkflow(ctx, sdkclient.StartWorkflowOptions{TaskQueue: tv1.TaskQueue().GetName()}, "wf") + run, err := env.SdkClient().ExecuteWorkflow(env.Context(), sdkclient.StartWorkflowOptions{TaskQueue: tv1.TaskQueue().GetName()}, "wf") s.NoError(err) runs[i] = run } for i := range numTests { var out string - s.NoError(runs[i].Get(ctx, &out)) + s.NoError(runs[i].Get(env.Context(), &out)) counter[out]++ } @@ -1303,63 +1266,65 @@ func (s *Versioning3Suite) TestTransitionFromWft_NoSticky_ToUnversioned() { } func (s *Versioning3Suite) testTransitionFromWft(sticky bool, toUnversioned bool) { + env := s.setupEnv() + // Wf runs one WFT and one AT on d1, then the second WFT is redirected to d2 and // transitions the wf with it. - tv1 := testvars.New(s).WithBuildIDNumber(1) + tv1 := env.Tv().WithBuildIDNumber(1) tv2 := tv1.WithBuildIDNumber(2) if sticky { - s.warmUpSticky(tv1) + s.warmUpSticky(env, tv1) } - s.updateTaskQueueDeploymentDataWithRoutingConfig(tv1, &deploymentpb.RoutingConfig{ + s.updateTaskQueueDeploymentDataWithRoutingConfig(env, tv1, &deploymentpb.RoutingConfig{ CurrentDeploymentVersion: worker_versioning.ExternalWorkerDeploymentVersionFromStringV31(tv1.DeploymentVersionString()), CurrentVersionChangedTime: timestamp.TimePtr(time.Now()), RevisionNumber: 1, }, map[string]*deploymentspb.WorkerDeploymentVersionData{tv1.DeploymentVersion().GetBuildId(): { Status: enumspb.WORKER_DEPLOYMENT_VERSION_STATUS_CURRENT, }}, []string{}, tqTypeWf, tqTypeAct) - runID := s.startWorkflow(tv1, nil) + runID := s.startWorkflow(env, tv1, nil) - s.pollWftAndHandle(tv1, false, nil, + s.pollWftAndHandle(env, tv1, false, nil, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) - s.verifyWorkflowVersioning(s.Assertions, tv1, vbUnspecified, nil, nil, tv1.DeploymentVersionTransition()) + s.verifyWorkflowVersioning(env, s.Assertions, tv1, vbUnspecified, nil, nil, tv1.DeploymentVersionTransition()) return respondWftWithActivities(tv1, tv1, sticky, vbUnpinned, "5"), nil }) - s.verifyWorkflowVersioning(s.Assertions, tv1, vbUnpinned, tv1.Deployment(), nil, nil) - s.verifyVersioningSAs(tv1, vbUnpinned, enumspb.WORKFLOW_EXECUTION_STATUS_RUNNING, tv1) + s.verifyWorkflowVersioning(env, s.Assertions, tv1, vbUnpinned, tv1.Deployment(), nil, nil) + s.verifyVersioningSAs(env, tv1, vbUnpinned, enumspb.WORKFLOW_EXECUTION_STATUS_RUNNING, tv1) if sticky { - s.verifyWorkflowStickyQueue(tv1.WithRunID(runID)) + s.verifyWorkflowStickyQueue(env, tv1.WithRunID(runID)) } - s.pollActivityAndHandle(tv1, nil, + s.pollActivityAndHandle(env, tv1, nil, func(task *workflowservice.PollActivityTaskQueueResponse) (*workflowservice.RespondActivityTaskCompletedRequest, error) { s.NotNil(task) return respondActivity(), nil }) - s.verifyWorkflowVersioning(s.Assertions, tv1, vbUnpinned, tv1.Deployment(), nil, nil) + s.verifyWorkflowVersioning(env, s.Assertions, tv1, vbUnpinned, tv1.Deployment(), nil, nil) if toUnversioned { // unset A as current - s.updateTaskQueueDeploymentDataWithRoutingConfig(tv1, &deploymentpb.RoutingConfig{ + s.updateTaskQueueDeploymentDataWithRoutingConfig(env, tv1, &deploymentpb.RoutingConfig{ CurrentDeploymentVersion: nil, CurrentVersionChangedTime: timestamp.TimePtr(time.Now()), RevisionNumber: 2, }, map[string]*deploymentspb.WorkerDeploymentVersionData{}, []string{}, tqTypeWf, tqTypeAct) - s.unversionedPollWftAndHandle(tv1, false, nil, + s.unversionedPollWftAndHandle(env, tv1, false, nil, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) - s.verifyWorkflowVersioning(s.Assertions, tv1, vbUnpinned, tv1.Deployment(), nil, &workflowpb.DeploymentVersionTransition{Version: "__unversioned__"}) + s.verifyWorkflowVersioning(env, s.Assertions, tv1, vbUnpinned, tv1.Deployment(), nil, &workflowpb.DeploymentVersionTransition{Version: "__unversioned__"}) return respondCompleteWorkflowUnversioned(tv1), nil }) - s.verifyWorkflowVersioning(s.Assertions, tv1, vbUnspecified, nil, nil, nil) - s.verifyVersioningSAs(tv1, vbUnspecified, enumspb.WORKFLOW_EXECUTION_STATUS_COMPLETED, tv1) + s.verifyWorkflowVersioning(env, s.Assertions, tv1, vbUnspecified, nil, nil, nil) + s.verifyVersioningSAs(env, tv1, vbUnspecified, enumspb.WORKFLOW_EXECUTION_STATUS_COMPLETED, tv1) } else { // Set B as the current deployment - s.updateTaskQueueDeploymentDataWithRoutingConfig(tv2, &deploymentpb.RoutingConfig{ + s.updateTaskQueueDeploymentDataWithRoutingConfig(env, tv2, &deploymentpb.RoutingConfig{ CurrentDeploymentVersion: worker_versioning.ExternalWorkerDeploymentVersionFromStringV31(tv2.DeploymentVersionString()), CurrentVersionChangedTime: timestamp.TimePtr(time.Now()), RevisionNumber: 2, @@ -1369,14 +1334,14 @@ func (s *Versioning3Suite) testTransitionFromWft(sticky bool, toUnversioned bool Status: enumspb.WORKER_DEPLOYMENT_VERSION_STATUS_DRAINING, }}, []string{}, tqTypeWf, tqTypeAct) - s.pollWftAndHandle(tv2, false, nil, + s.pollWftAndHandle(env, tv2, false, nil, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) - s.verifyWorkflowVersioning(s.Assertions, tv1, vbUnpinned, tv1.Deployment(), nil, tv2.DeploymentVersionTransition()) + s.verifyWorkflowVersioning(env, s.Assertions, tv1, vbUnpinned, tv1.Deployment(), nil, tv2.DeploymentVersionTransition()) return respondCompleteWorkflow(tv2, vbUnpinned), nil }) - s.verifyWorkflowVersioning(s.Assertions, tv2, vbUnpinned, tv2.Deployment(), nil, nil) - s.verifyVersioningSAs(tv2, vbUnpinned, enumspb.WORKFLOW_EXECUTION_STATUS_COMPLETED, tv1, tv2) + s.verifyWorkflowVersioning(env, s.Assertions, tv2, vbUnpinned, tv2.Deployment(), nil, nil) + s.verifyVersioningSAs(env, tv2, vbUnpinned, enumspb.WORKFLOW_EXECUTION_STATUS_COMPLETED, tv1, tv2) } } @@ -1397,6 +1362,8 @@ func (s *Versioning3Suite) TestDoubleTransitionFromUnversioned_WithSignal() { } func (s *Versioning3Suite) testDoubleTransition(unversionedSrc bool, signal bool) { + env := s.setupEnv() + // Scenario: make sure a transition on top of another transition works properly. Steps: // 1. Start a workflow, process first task with sourceV (which is v1 or unversioned depending on passed `unversionedSrc`). // The first wft schedules an activity. @@ -1411,7 +1378,7 @@ func (s *Versioning3Suite) testDoubleTransition(unversionedSrc bool, signal bool // 8. Not that the transition is complete, the activity should be able to go to sourceV poller. // 9. Set v2 as current again and ensure the wf can complete on it. - tv1 := testvars.New(s).WithBuildIDNumber(1) + tv1 := env.Tv().WithBuildIDNumber(1) tv2 := tv1.WithBuildIDNumber(2) sourceV := tv1.Deployment() sourceVB := vbUnpinned @@ -1422,11 +1389,11 @@ func (s *Versioning3Suite) testDoubleTransition(unversionedSrc bool, signal bool sourceTransition = &workflowpb.DeploymentVersionTransition{Version: "__unversioned__"} } - runId := s.startWorkflow(tv1, nil) + runID := s.startWorkflow(env, tv1, nil) if !unversionedSrc { // sourceV is v1, set current version to it - s.updateTaskQueueDeploymentDataWithRoutingConfig(tv1, &deploymentpb.RoutingConfig{ + s.updateTaskQueueDeploymentDataWithRoutingConfig(env, tv1, &deploymentpb.RoutingConfig{ CurrentDeploymentVersion: worker_versioning.ExternalWorkerDeploymentVersionFromStringV31(tv1.DeploymentVersionString()), CurrentVersionChangedTime: timestamp.TimePtr(time.Now()), RevisionNumber: 1, @@ -1435,21 +1402,21 @@ func (s *Versioning3Suite) testDoubleTransition(unversionedSrc bool, signal bool }}, []string{}, tqTypeWf, tqTypeAct) } - s.doPollWftAndHandle(tv1, !unversionedSrc, false, nil, + s.doPollWftAndHandle(env, tv1, !unversionedSrc, false, nil, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) return respondWftWithActivities(tv1, tv1, false, sourceVB, "5"), nil }) - s.verifyWorkflowVersioning(s.Assertions, tv1, sourceVB, sourceV, nil, nil) + s.verifyWorkflowVersioning(env, s.Assertions, tv1, sourceVB, sourceV, nil, nil) if signal { // Send a signal so a wf task is scheduled before we poll the activity - err := s.SdkClient().SignalWorkflow(context.Background(), tv1.WorkflowID(), runId, "test-signal", nil) + err := env.SdkClient().SignalWorkflow(env.Context(), tv1.WorkflowID(), runID, "test-signal", nil) s.NoError(err) } // set current version to v2 - s.updateTaskQueueDeploymentDataWithRoutingConfig(tv2, &deploymentpb.RoutingConfig{ + s.updateTaskQueueDeploymentDataWithRoutingConfig(env, tv2, &deploymentpb.RoutingConfig{ CurrentDeploymentVersion: worker_versioning.ExternalWorkerDeploymentVersionFromStringV31(tv2.DeploymentVersionString()), CurrentVersionChangedTime: timestamp.TimePtr(time.Now()), RevisionNumber: 2, @@ -1459,12 +1426,12 @@ func (s *Versioning3Suite) testDoubleTransition(unversionedSrc bool, signal bool Status: enumspb.WORKER_DEPLOYMENT_VERSION_STATUS_DRAINING, }}, []string{}, tqTypeWf, tqTypeAct) // poll activity from v2 worker, this should start a transition but should not immediately start the activity. - go s.idlePollActivity(context.Background(), tv2, true, time.Minute, "v2 worker should not receive the activity") + go s.idlePollActivity(env, tv2, true, time.Minute, "v2 worker should not receive the activity") s.EventuallyWithT(func(t *assert.CollectT) { - dwf, err := s.FrontendClient().DescribeWorkflowExecution( - context.Background(), &workflowservice.DescribeWorkflowExecutionRequest{ - Namespace: s.Namespace().String(), + dwf, err := env.FrontendClient().DescribeWorkflowExecution( + env.Context(), &workflowservice.DescribeWorkflowExecutionRequest{ + Namespace: env.Namespace().String(), Execution: &commonpb.WorkflowExecution{ WorkflowId: tv1.WorkflowID(), }, @@ -1477,13 +1444,13 @@ func (s *Versioning3Suite) testDoubleTransition(unversionedSrc bool, signal bool // Back to sourceV if unversionedSrc { - s.updateTaskQueueDeploymentDataWithRoutingConfig(tv2, &deploymentpb.RoutingConfig{ + s.updateTaskQueueDeploymentDataWithRoutingConfig(env, tv2, &deploymentpb.RoutingConfig{ CurrentDeploymentVersion: nil, CurrentVersionChangedTime: timestamp.TimePtr(time.Now()), RevisionNumber: 3, }, map[string]*deploymentspb.WorkerDeploymentVersionData{}, []string{}, tqTypeWf, tqTypeAct) } else { - s.updateTaskQueueDeploymentDataWithRoutingConfig(tv1, &deploymentpb.RoutingConfig{ + s.updateTaskQueueDeploymentDataWithRoutingConfig(env, tv1, &deploymentpb.RoutingConfig{ CurrentDeploymentVersion: worker_versioning.ExternalWorkerDeploymentVersionFromStringV31(tv1.DeploymentVersionString()), CurrentVersionChangedTime: timestamp.TimePtr(time.Now()), RevisionNumber: 3, @@ -1496,24 +1463,24 @@ func (s *Versioning3Suite) testDoubleTransition(unversionedSrc bool, signal bool // Now poll for wf task from sourceV while there is a transition to v2 // this poller should clear the transition and complete the wft. - s.doPollWftAndHandle(tv1, !unversionedSrc, false, nil, + s.doPollWftAndHandle(env, tv1, !unversionedSrc, false, nil, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) - s.verifyWorkflowVersioning(s.Assertions, tv1, sourceVB, sourceV, nil, sourceTransition) + s.verifyWorkflowVersioning(env, s.Assertions, tv1, sourceVB, sourceV, nil, sourceTransition) return respondEmptyWft(tv1, false, sourceVB), nil }) - s.verifyWorkflowVersioning(s.Assertions, tv1, sourceVB, sourceV, nil, nil) + s.verifyWorkflowVersioning(env, s.Assertions, tv1, sourceVB, sourceV, nil, nil) // Activity should be unblocked now to sourceV poller - s.doPollActivityAndHandle(tv1, !unversionedSrc, nil, + s.doPollActivityAndHandle(env, tv1, !unversionedSrc, nil, func(task *workflowservice.PollActivityTaskQueueResponse) (*workflowservice.RespondActivityTaskCompletedRequest, error) { s.NotNil(task) - s.verifyWorkflowVersioning(s.Assertions, tv1, sourceVB, sourceV, nil, nil) + s.verifyWorkflowVersioning(env, s.Assertions, tv1, sourceVB, sourceV, nil, nil) return respondActivity(), nil }) // Set v2 as the current version again - s.updateTaskQueueDeploymentDataWithRoutingConfig(tv2, &deploymentpb.RoutingConfig{ + s.updateTaskQueueDeploymentDataWithRoutingConfig(env, tv2, &deploymentpb.RoutingConfig{ CurrentDeploymentVersion: worker_versioning.ExternalWorkerDeploymentVersionFromStringV31(tv2.DeploymentVersionString()), CurrentVersionChangedTime: timestamp.TimePtr(time.Now()), RevisionNumber: 4, @@ -1523,29 +1490,27 @@ func (s *Versioning3Suite) testDoubleTransition(unversionedSrc bool, signal bool Status: enumspb.WORKER_DEPLOYMENT_VERSION_STATUS_DRAINING, }}, []string{}, tqTypeWf, tqTypeAct) - s.pollWftAndHandle(tv2, false, nil, + s.pollWftAndHandle(env, tv2, false, nil, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) - s.verifyWorkflowVersioning(s.Assertions, tv2, sourceVB, sourceV, nil, tv2.DeploymentVersionTransition()) + s.verifyWorkflowVersioning(env, s.Assertions, tv2, sourceVB, sourceV, nil, tv2.DeploymentVersionTransition()) return respondCompleteWorkflow(tv2, vbUnpinned), nil }) - s.verifyWorkflowVersioning(s.Assertions, tv2, vbUnpinned, tv2.Deployment(), nil, nil) + s.verifyWorkflowVersioning(env, s.Assertions, tv2, vbUnpinned, tv2.Deployment(), nil, nil) } func (s *Versioning3Suite) TestNexusTask_StaysOnCurrentDeployment() { - s.RunTestWithMatchingBehavior( - func() { - s.nexusTaskStaysOnCurrentDeployment() - }, - ) + s.runTestWithMatchingBehavior(func(env *testcore.TestEnv, s *Versioning3Suite) { + s.nexusTaskStaysOnCurrentDeployment(env) + }) } -func (s *Versioning3Suite) nexusTaskStaysOnCurrentDeployment() { - tv1 := testvars.New(s).WithBuildIDNumber(1) +func (s *Versioning3Suite) nexusTaskStaysOnCurrentDeployment(env *testcore.TestEnv) { + tv1 := env.Tv().WithBuildIDNumber(1) tv2 := tv1.WithBuildIDNumber(2) nexusRequest := &matchingservice.DispatchNexusTaskRequest{ - NamespaceId: s.NamespaceID().String(), + NamespaceId: env.NamespaceID().String(), TaskQueue: tv1.TaskQueue(), Request: &nexuspb.Request{ Header: map[string]string{ @@ -1557,7 +1522,7 @@ func (s *Versioning3Suite) nexusTaskStaysOnCurrentDeployment() { } // current deployment is -> tv1 - s.updateTaskQueueDeploymentDataWithRoutingConfig(tv1, &deploymentpb.RoutingConfig{ + s.updateTaskQueueDeploymentDataWithRoutingConfig(env, tv1, &deploymentpb.RoutingConfig{ CurrentDeploymentVersion: worker_versioning.ExternalWorkerDeploymentVersionFromStringV31(tv1.DeploymentVersionString()), CurrentVersionChangedTime: timestamp.TimePtr(time.Now()), RevisionNumber: 1, @@ -1566,10 +1531,10 @@ func (s *Versioning3Suite) nexusTaskStaysOnCurrentDeployment() { }}, []string{}, tqTypeNexus) // local poller with deployment A receives task - s.pollAndDispatchNexusTask(tv1, nexusRequest) + s.pollAndDispatchNexusTask(env, tv1, nexusRequest) // current deployment is now -> tv2 - s.updateTaskQueueDeploymentDataWithRoutingConfig(tv2, &deploymentpb.RoutingConfig{ + s.updateTaskQueueDeploymentDataWithRoutingConfig(env, tv2, &deploymentpb.RoutingConfig{ CurrentDeploymentVersion: worker_versioning.ExternalWorkerDeploymentVersionFromStringV31(tv2.DeploymentVersionString()), CurrentVersionChangedTime: timestamp.TimePtr(time.Now()), RevisionNumber: 2, @@ -1580,49 +1545,52 @@ func (s *Versioning3Suite) nexusTaskStaysOnCurrentDeployment() { }}, []string{}, tqTypeNexus) // Pollers of tv1 are there but should not get any task - go s.idlePollNexus(context.Background(), tv1, true, ver3MinPollTime, "nexus task should not go to the old deployment") + go s.idlePollNexus(env, env.Context(), tv1, true, ver3MinPollTime, "nexus task should not go to the old deployment") - s.pollAndDispatchNexusTask(tv2, nexusRequest) + s.pollAndDispatchNexusTask(env, tv2, nexusRequest) } func (s *Versioning3Suite) pollAndDispatchNexusTask( + env *testcore.TestEnv, tv *testvars.TestVars, nexusRequest *matchingservice.DispatchNexusTaskRequest, ) { - matchingClient := s.GetTestCluster().MatchingClient() + matchingClient := env.GetTestCluster().MatchingClient() nexusCompleted := make(chan any) - s.pollNexusTaskAndHandle(tv, false, nexusCompleted, + s.pollNexusTaskAndHandle(env, tv, false, nexusCompleted, func(task *workflowservice.PollNexusTaskQueueResponse) (*workflowservice.RespondNexusTaskCompletedRequest, error) { s.NotNil(task) return &workflowservice.RespondNexusTaskCompletedRequest{}, nil // response object gets filled during processing }) - _, err := matchingClient.DispatchNexusTask(context.Background(), nexusRequest) + _, err := matchingClient.DispatchNexusTask(env.Context(), nexusRequest) s.NoError(err) <-nexusCompleted } func (s *Versioning3Suite) TestEagerActivity() { + env := s.setupEnv( + testcore.WithDynamicConfig(dynamicconfig.EnableActivityEagerExecution, true), + ) // The first WFT asks for an activity to starts and get it eagerly in the WFT completion // response. The activity is processed without issues and wf completes. - s.OverrideDynamicConfig(dynamicconfig.EnableActivityEagerExecution, true) - tv := testvars.New(s) + tv := env.Tv() - s.updateTaskQueueDeploymentDataWithRoutingConfig(tv, &deploymentpb.RoutingConfig{ + s.updateTaskQueueDeploymentDataWithRoutingConfig(env, tv, &deploymentpb.RoutingConfig{ CurrentDeploymentVersion: worker_versioning.ExternalWorkerDeploymentVersionFromStringV31(tv.DeploymentVersionString()), CurrentVersionChangedTime: timestamp.TimePtr(time.Now()), RevisionNumber: 1, }, map[string]*deploymentspb.WorkerDeploymentVersionData{tv.DeploymentVersion().GetBuildId(): { Status: enumspb.WORKER_DEPLOYMENT_VERSION_STATUS_CURRENT, }}, []string{}, tqTypeWf, tqTypeAct) - s.startWorkflow(tv, nil) + s.startWorkflow(env, tv, nil) - poller, resp := s.pollWftAndHandle(tv, false, nil, + poller, resp := s.pollWftAndHandle(env, tv, false, nil, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) - s.verifyWorkflowVersioning(s.Assertions, tv, vbUnspecified, nil, nil, tv.DeploymentVersionTransition()) + s.verifyWorkflowVersioning(env, s.Assertions, tv, vbUnspecified, nil, nil, tv.DeploymentVersionTransition()) resp := respondWftWithActivities(tv, tv, true, vbUnpinned, "5") resp.Commands[0].GetScheduleActivityTaskCommandAttributes().RequestEagerExecution = true resp.Commands[0].GetScheduleActivityTaskCommandAttributes().Priority = &commonpb.Priority{ @@ -1631,7 +1599,7 @@ func (s *Versioning3Suite) TestEagerActivity() { } return resp, nil }) - s.verifyWorkflowVersioning(s.Assertions, tv, vbUnpinned, tv.Deployment(), nil, nil) + s.verifyWorkflowVersioning(env, s.Assertions, tv, vbUnpinned, tv.Deployment(), nil, nil) s.NotEmpty(resp.GetActivityTasks()) s.Equal("fairness-key", resp.GetActivityTasks()[0].GetPriority().GetFairnessKey()) @@ -1643,14 +1611,14 @@ func (s *Versioning3Suite) TestEagerActivity() { return respondActivity(), nil }) s.NoError(err) - s.verifyWorkflowVersioning(s.Assertions, tv, vbUnpinned, tv.Deployment(), nil, nil) + s.verifyWorkflowVersioning(env, s.Assertions, tv, vbUnpinned, tv.Deployment(), nil, nil) - s.pollWftAndHandle(tv, false, nil, + s.pollWftAndHandle(env, tv, false, nil, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) return respondCompleteWorkflow(tv, vbUnpinned), nil }) - s.verifyWorkflowVersioning(s.Assertions, tv, vbUnpinned, tv.Deployment(), nil, nil) + s.verifyWorkflowVersioning(env, s.Assertions, tv, vbUnpinned, tv.Deployment(), nil, nil) } func (s *Versioning3Suite) TestTransitionFromActivity_Sticky() { @@ -1662,6 +1630,8 @@ func (s *Versioning3Suite) TestTransitionFromActivity_NoSticky() { } func (s *Versioning3Suite) testTransitionFromActivity(sticky bool) { + env := s.setupEnv() + // The workflow runs one WFT on d1 which schedules four activities, then: // 1. The first and second activities starts on d1 // 2. Current deployment becomes d2 @@ -1675,33 +1645,30 @@ func (s *Versioning3Suite) testTransitionFromActivity(sticky bool) { // 8. WFT completes and the transition completes. // 9. All the 3 remaining activities are now dispatched and completed. - ctx, cancel := context.WithTimeout(context.Background(), 90*time.Second) - defer cancel() - - tv1 := testvars.New(s).WithBuildIDNumber(1) + tv1 := env.Tv().WithBuildIDNumber(1) tv2 := tv1.WithBuildIDNumber(2) if sticky { - s.warmUpSticky(tv1) + s.warmUpSticky(env, tv1) } - s.updateTaskQueueDeploymentDataWithRoutingConfig(tv1, &deploymentpb.RoutingConfig{ + s.updateTaskQueueDeploymentDataWithRoutingConfig(env, tv1, &deploymentpb.RoutingConfig{ CurrentDeploymentVersion: worker_versioning.ExternalWorkerDeploymentVersionFromStringV31(tv1.DeploymentVersionString()), CurrentVersionChangedTime: timestamp.TimePtr(time.Now()), RevisionNumber: 1, }, map[string]*deploymentspb.WorkerDeploymentVersionData{tv1.DeploymentVersion().GetBuildId(): { Status: enumspb.WORKER_DEPLOYMENT_VERSION_STATUS_CURRENT, }}, []string{}, tqTypeWf, tqTypeAct) - runID := s.startWorkflow(tv1, nil) + runID := s.startWorkflow(env, tv1, nil) - s.pollWftAndHandle(tv1, false, nil, + s.pollWftAndHandle(env, tv1, false, nil, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) - s.verifyWorkflowVersioning(s.Assertions, tv1, vbUnspecified, nil, nil, tv1.DeploymentVersionTransition()) + s.verifyWorkflowVersioning(env, s.Assertions, tv1, vbUnspecified, nil, nil, tv1.DeploymentVersionTransition()) return respondWftWithActivities(tv1, tv1, sticky, vbUnpinned, "5", "6", "7", "8"), nil }) - s.verifyWorkflowVersioning(s.Assertions, tv1, vbUnpinned, tv1.Deployment(), nil, nil) + s.verifyWorkflowVersioning(env, s.Assertions, tv1, vbUnpinned, tv1.Deployment(), nil, nil) if sticky { - s.verifyWorkflowStickyQueue(tv1.WithRunID(runID)) + s.verifyWorkflowStickyQueue(env, tv1.WithRunID(runID)) } transitionCompleted := atomic.Bool{} @@ -1713,44 +1680,44 @@ func (s *Versioning3Suite) testTransitionFromActivity(sticky bool) { act2To4Completed := make(chan struct{}) // 1. Start 1st and 2nd activities - s.pollActivityAndHandle(tv1, act1Completed, + s.pollActivityAndHandle(env, tv1, act1Completed, func(task *workflowservice.PollActivityTaskQueueResponse) (*workflowservice.RespondActivityTaskCompletedRequest, error) { s.NotNil(task) - s.Logger.Info(fmt.Sprintf("Activity 1 started ID: %s", task.ActivityId)) + env.Logger.Info(fmt.Sprintf("Activity 1 started ID: %s", task.ActivityId)) close(act1Started) // block until the transition WFT starts select { case <-transitionStarted: - case <-ctx.Done(): - return nil, fmt.Errorf("context expired waiting for transitionStarted in act1: %w", ctx.Err()) + case <-env.Context().Done(): + return nil, fmt.Errorf("context expired waiting for transitionStarted in act1: %w", env.Context().Err()) } // 6. the 1st act completes during transition - s.Logger.Info(fmt.Sprintf("Activity 1 completed ID: %s", task.ActivityId)) + env.Logger.Info(fmt.Sprintf("Activity 1 completed ID: %s", task.ActivityId)) return respondActivity(), nil }) - s.WaitForChannel(ctx, act1Started) - s.pollActivityAndHandle(tv1, act2Failed, + env.WaitForChannel(env.Context(), act1Started) + s.pollActivityAndHandle(env, tv1, act2Failed, func(task *workflowservice.PollActivityTaskQueueResponse) (*workflowservice.RespondActivityTaskCompletedRequest, error) { s.NotNil(task) - s.Logger.Info(fmt.Sprintf("Activity 2 started ID: %s", task.ActivityId)) + env.Logger.Info(fmt.Sprintf("Activity 2 started ID: %s", task.ActivityId)) close(act2Started) // block until the transition WFT starts select { case <-transitionStarted: - case <-ctx.Done(): - return nil, fmt.Errorf("context expired waiting for transitionStarted in act2: %w", ctx.Err()) + case <-env.Context().Done(): + return nil, fmt.Errorf("context expired waiting for transitionStarted in act2: %w", env.Context().Err()) } // 7. 2nd activity fails. Respond with error so it is retried. - s.Logger.Info(fmt.Sprintf("Activity 2 failed ID: %s", task.ActivityId)) + env.Logger.Info(fmt.Sprintf("Activity 2 failed ID: %s", task.ActivityId)) return nil, errors.New("intentional activity failure") }) - s.WaitForChannel(ctx, act2Started) - s.verifyWorkflowVersioning(s.Assertions, tv1, vbUnpinned, tv1.Deployment(), nil, nil) + env.WaitForChannel(env.Context(), act2Started) + s.verifyWorkflowVersioning(env, s.Assertions, tv1, vbUnpinned, tv1.Deployment(), nil, nil) // 2. Set d2 as the current deployment - s.updateTaskQueueDeploymentDataWithRoutingConfig(tv2, &deploymentpb.RoutingConfig{ + s.updateTaskQueueDeploymentDataWithRoutingConfig(env, tv2, &deploymentpb.RoutingConfig{ CurrentDeploymentVersion: worker_versioning.ExternalWorkerDeploymentVersionFromStringV31(tv2.DeploymentVersionString()), CurrentVersionChangedTime: timestamp.TimePtr(time.Now()), RevisionNumber: 2, @@ -1766,61 +1733,75 @@ func (s *Versioning3Suite) testTransitionFromActivity(sticky bool) { time.Sleep(time.Millisecond * 200) //nolint:forbidigo // Pollers of d1 are there, but should not get any task - go s.idlePollActivity(ctx, tv1, true, ver3MinPollTime, "activities should not go to the old deployment") + go s.idlePollActivity(env, tv1, true, ver3MinPollTime, "activities should not go to the old deployment") + act2To4Err := make(chan error, 1) go func() { + defer close(act2To4Completed) for i := 2; i <= 4; i++ { // 3-4. The new d2 poller should trigger the third activity to be redirected, but the activity should // not start until transition completes in the next wft. // Repeating the handler so it processes all the three activities - s.pollActivityAndHandle(tv2, nil, + err := s.pollActivityAndHandleErr(env, tv2, func(task *workflowservice.PollActivityTaskQueueResponse) (*workflowservice.RespondActivityTaskCompletedRequest, error) { // Activity should not start until the transition is completed - s.True(transitionCompleted.Load()) - s.NotNil(task) - s.Logger.Info(fmt.Sprintf("Remaining activity completed ID: %s", task.ActivityId)) + if !transitionCompleted.Load() { + return nil, errors.New("activity started before transition completed") + } + if task == nil { + return nil, errors.New("missing activity task") + } + env.Logger.Info(fmt.Sprintf("Remaining activity completed ID: %s", task.ActivityId)) return respondActivity(), nil }) + if err != nil { + act2To4Err <- err + return + } } - close(act2To4Completed) }() // 5. The transition should create a new WFT to be sent to d2. - s.pollWftAndHandle(tv2, false, nil, + s.pollWftAndHandle(env, tv2, false, nil, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) - s.verifyWorkflowVersioning(s.Assertions, tv1, vbUnpinned, tv1.Deployment(), nil, tv2.DeploymentVersionTransition()) + s.verifyWorkflowVersioning(env, s.Assertions, tv1, vbUnpinned, tv1.Deployment(), nil, tv2.DeploymentVersionTransition()) close(transitionStarted) - s.Logger.Info("Transition wft started") + env.Logger.Info("Transition wft started") // 8. Complete the transition after act1 completes and act2's first attempt fails. select { case <-act1Completed: - case <-ctx.Done(): + case <-env.Context().Done(): s.FailNow("context expired waiting for act1 to complete") } select { case <-act2Failed: - case <-ctx.Done(): + case <-env.Context().Done(): s.FailNow("context expired waiting for act2 to fail") } transitionCompleted.Store(true) - s.Logger.Info("Transition wft completed") + env.Logger.Info("Transition wft completed") return respondEmptyWft(tv2, sticky, vbUnpinned), nil }) - s.verifyWorkflowVersioning(s.Assertions, tv2, vbUnpinned, tv2.Deployment(), nil, nil) + s.verifyWorkflowVersioning(env, s.Assertions, tv2, vbUnpinned, tv2.Deployment(), nil, nil) if sticky { - s.verifyWorkflowStickyQueue(tv2) + s.verifyWorkflowStickyQueue(env, tv2) } // 9. Now all activities should complete. - s.WaitForChannel(ctx, act2To4Completed) - s.pollWftAndHandle(tv2, sticky, nil, + env.WaitForChannel(env.Context(), act2To4Completed) + select { + case err := <-act2To4Err: + s.NoError(err) + default: + } + s.pollWftAndHandle(env, tv2, sticky, nil, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) - s.Logger.Info("Final wft completed") + env.Logger.Info("Final wft completed") return respondCompleteWorkflow(tv2, vbUnpinned), nil }) - s.verifyWorkflowVersioning(s.Assertions, tv2, vbUnpinned, tv2.Deployment(), nil, nil) + s.verifyWorkflowVersioning(env, s.Assertions, tv2, vbUnpinned, tv2.Deployment(), nil, nil) } func (s *Versioning3Suite) TestIndependentVersionedActivity_Pinned() { @@ -1840,17 +1821,19 @@ func (s *Versioning3Suite) TestIndependentUnversionedActivity_Unpinned() { } func (s *Versioning3Suite) testIndependentActivity(behavior enumspb.VersioningBehavior, unversionedActivity bool) { + env := s.setupEnv() + // This test starts a wf on wf-series. The workflow runs an activity that is sent to act-tq with // workers on a different deployment series, act-series. We make sure that the activity is // dispatched and processed properly without affecting versioning of the workflow. Note that it // is not required for independent activities to use a different TQ name but in here we test the // more common case where the TQ name is different. - tvWf := testvars.New(s).WithDeploymentSeriesNumber(1) - tvAct := testvars.New(s).WithDeploymentSeriesNumber(2).WithTaskQueueNumber(2) + tvWf := env.Tv().WithDeploymentSeriesNumber(1) + tvAct := env.Tv().WithDeploymentSeriesNumber(2).WithTaskQueueNumber(2) // Set current deployment for each TQ - s.updateTaskQueueDeploymentDataWithRoutingConfig(tvWf, &deploymentpb.RoutingConfig{ + s.updateTaskQueueDeploymentDataWithRoutingConfig(env, tvWf, &deploymentpb.RoutingConfig{ CurrentDeploymentVersion: worker_versioning.ExternalWorkerDeploymentVersionFromStringV31(tvWf.DeploymentVersionString()), CurrentVersionChangedTime: timestamp.TimePtr(time.Now()), RevisionNumber: 1, @@ -1860,7 +1843,7 @@ func (s *Versioning3Suite) testIndependentActivity(behavior enumspb.VersioningBe if !unversionedActivity { // Different deployment here for the activity TQ. - s.updateTaskQueueDeploymentDataWithRoutingConfig(tvAct, &deploymentpb.RoutingConfig{ + s.updateTaskQueueDeploymentDataWithRoutingConfig(env, tvAct, &deploymentpb.RoutingConfig{ CurrentDeploymentVersion: worker_versioning.ExternalWorkerDeploymentVersionFromStringV31(tvAct.DeploymentVersionString()), CurrentVersionChangedTime: timestamp.TimePtr(time.Now()), RevisionNumber: 1, @@ -1869,60 +1852,62 @@ func (s *Versioning3Suite) testIndependentActivity(behavior enumspb.VersioningBe }}, []string{}, tqTypeAct) } - s.startWorkflow(tvWf, nil) + s.startWorkflow(env, tvWf, nil) - s.pollWftAndHandle(tvWf, false, nil, + s.pollWftAndHandle(env, tvWf, false, nil, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) - s.verifyWorkflowVersioning(s.Assertions, tvWf, vbUnspecified, nil, nil, tvWf.DeploymentVersionTransition()) - s.Logger.Info("First wf task completed") + s.verifyWorkflowVersioning(env, s.Assertions, tvWf, vbUnspecified, nil, nil, tvWf.DeploymentVersionTransition()) + env.Logger.Info("First wf task completed") return respondWftWithActivities(tvWf, tvAct, false, behavior, "5"), nil }) - s.verifyWorkflowVersioning(s.Assertions, tvWf, behavior, tvWf.Deployment(), nil, nil) + s.verifyWorkflowVersioning(env, s.Assertions, tvWf, behavior, tvWf.Deployment(), nil, nil) if unversionedActivity { - s.unversionedPollActivityAndHandle(tvAct, nil, + s.unversionedPollActivityAndHandle(env, tvAct, nil, func(task *workflowservice.PollActivityTaskQueueResponse) (*workflowservice.RespondActivityTaskCompletedRequest, error) { s.NotNil(task) - s.Logger.Info("Activity completed") + env.Logger.Info("Activity completed") return respondActivity(), nil }) } else { - s.pollActivityAndHandle(tvAct, nil, + s.pollActivityAndHandle(env, tvAct, nil, func(task *workflowservice.PollActivityTaskQueueResponse) (*workflowservice.RespondActivityTaskCompletedRequest, error) { s.NotNil(task) - s.Logger.Info("Activity completed") + env.Logger.Info("Activity completed") return respondActivity(), nil }) } - s.verifyWorkflowVersioning(s.Assertions, tvWf, behavior, tvWf.Deployment(), nil, nil) + s.verifyWorkflowVersioning(env, s.Assertions, tvWf, behavior, tvWf.Deployment(), nil, nil) - s.pollWftAndHandle(tvWf, false, nil, + s.pollWftAndHandle(env, tvWf, false, nil, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) return respondCompleteWorkflow(tvWf, behavior), nil }) - s.verifyWorkflowVersioning(s.Assertions, tvWf, behavior, tvWf.Deployment(), nil, nil) + s.verifyWorkflowVersioning(env, s.Assertions, tvWf, behavior, tvWf.Deployment(), nil, nil) } func (s *Versioning3Suite) TestChildWorkflowInheritance_PinnedParent() { - s.testChildWorkflowInheritance_ExpectInherit(false, false, vbPinned) + s.testChildWorkflowInheritanceExpectInherit(false, false, vbPinned) } func (s *Versioning3Suite) TestChildWorkflowInheritance_ParentPinnedByOverride() { - s.testChildWorkflowInheritance_ExpectInherit(false, true, vbUnpinned) + s.testChildWorkflowInheritanceExpectInherit(false, true, vbUnpinned) } func (s *Versioning3Suite) TestChildWorkflowInheritance_CrossTQ_Inherit() { // the ExpectInherit helper polls on the child's task queue with the parent's version, // so we expect the version to be inherited - s.testChildWorkflowInheritance_ExpectInherit(true, false, vbPinned) + s.testChildWorkflowInheritanceExpectInherit(true, false, vbPinned) } -func (s *Versioning3Suite) testChildWorkflowInheritance_ExpectInherit(crossTq bool, withOverride bool, parentRegistrationBehavior enumspb.VersioningBehavior) { +func (s *Versioning3Suite) testChildWorkflowInheritanceExpectInherit(crossTq bool, withOverride bool, parentRegistrationBehavior enumspb.VersioningBehavior) { + env := s.setupEnv() + // Child wf of a pinned parent starts on the parents pinned version. - tv1 := testvars.New(s).WithBuildIDNumber(1).WithWorkflowIDNumber(1) + tv1 := env.Tv().WithBuildIDNumber(1).WithWorkflowIDNumber(1) tv2 := tv1.WithBuildIDNumber(2) tv1Child := tv1.WithWorkflowIDNumber(2) if crossTq { @@ -1942,14 +1927,11 @@ func (s *Versioning3Suite) testChildWorkflowInheritance_ExpectInherit(crossTq bo sdkParentBehavior = workflow.VersioningBehaviorAutoUpgrade } - ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) - defer cancel() - wfStarted := make(chan struct{}, 1) currentChanged := make(chan struct{}, 1) childv1 := func(ctx workflow.Context) (string, error) { - s.verifyWorkflowVersioning(s.Assertions, tv1Child, vbPinned, tv1Child.Deployment(), override, nil) + s.verifyWorkflowVersioning(env, s.Assertions, tv1Child, vbPinned, tv1Child.Deployment(), override, nil) return "v1", nil } wf1 := func(ctx workflow.Context) (string, error) { @@ -1965,12 +1947,12 @@ func (s *Versioning3Suite) testChildWorkflowInheritance_ExpectInherit(crossTq bo var val1 string s.NoError(fut1.Get(ctx, &val1)) - s.verifyWorkflowVersioning(s.Assertions, tv1, parentRegistrationBehavior, tv1.Deployment(), override, nil) + s.verifyWorkflowVersioning(env, s.Assertions, tv1, parentRegistrationBehavior, tv1.Deployment(), override, nil) return val1, nil } if crossTq { - w1xtq := worker.New(s.SdkClient(), tv1Child.TaskQueue().GetName(), worker.Options{ + w1xtq := worker.New(env.SdkClient(), tv1Child.TaskQueue().GetName(), worker.Options{ DeploymentOptions: worker.DeploymentOptions{ Version: tv1Child.SDKDeploymentVersion(), UseVersioning: true, @@ -1983,7 +1965,7 @@ func (s *Versioning3Suite) testChildWorkflowInheritance_ExpectInherit(crossTq bo defer w1xtq.Stop() } - w1 := worker.New(s.SdkClient(), tv1.TaskQueue().GetName(), worker.Options{ + w1 := worker.New(env.SdkClient(), tv1.TaskQueue().GetName(), worker.Options{ DeploymentOptions: worker.DeploymentOptions{ Version: tv1.SDKDeploymentVersion(), UseVersioning: true, @@ -1999,7 +1981,7 @@ func (s *Versioning3Suite) testChildWorkflowInheritance_ExpectInherit(crossTq bo defer w1.Stop() // v1 is current for both parent and child - s.setCurrentDeployment(tv1) + s.setCurrentDeployment(env, tv1) startOpts := sdkclient.StartWorkflowOptions{ ID: tv1.WorkflowID(), @@ -2012,14 +1994,14 @@ func (s *Versioning3Suite) testChildWorkflowInheritance_ExpectInherit(crossTq bo Version: tv1.SDKDeploymentVersion(), } } - run, err := s.SdkClient().ExecuteWorkflow(ctx, startOpts, "wf") + run, err := env.SdkClient().ExecuteWorkflow(env.Context(), startOpts, "wf") s.NoError(err) // wait for it to start on v1 - s.WaitForChannel(ctx, wfStarted) + env.WaitForChannel(env.Context(), wfStarted) close(wfStarted) // force panic if replayed // make v2 current for both parent and child and unblock the wf to start the child - s.updateTaskQueueDeploymentDataWithRoutingConfig(tv2, &deploymentpb.RoutingConfig{ + s.updateTaskQueueDeploymentDataWithRoutingConfig(env, tv2, &deploymentpb.RoutingConfig{ CurrentDeploymentVersion: worker_versioning.ExternalWorkerDeploymentVersionFromStringV31(tv2.DeploymentVersionString()), CurrentVersionChangedTime: timestamp.TimePtr(time.Now()), RevisionNumber: 2, @@ -2027,7 +2009,7 @@ func (s *Versioning3Suite) testChildWorkflowInheritance_ExpectInherit(crossTq bo Status: enumspb.WORKER_DEPLOYMENT_VERSION_STATUS_CURRENT, }}, []string{}, tqTypeWf) if crossTq { - s.updateTaskQueueDeploymentDataWithRoutingConfig(tv2Child, &deploymentpb.RoutingConfig{ + s.updateTaskQueueDeploymentDataWithRoutingConfig(env, tv2Child, &deploymentpb.RoutingConfig{ CurrentDeploymentVersion: worker_versioning.ExternalWorkerDeploymentVersionFromStringV31(tv2Child.DeploymentVersionString()), CurrentVersionChangedTime: timestamp.TimePtr(time.Now()), RevisionNumber: 2, @@ -2038,26 +2020,28 @@ func (s *Versioning3Suite) testChildWorkflowInheritance_ExpectInherit(crossTq bo currentChanged <- struct{}{} var out string - s.NoError(run.Get(ctx, &out)) + s.NoError(run.Get(env.Context(), &out)) s.Equal("v1", out) } func (s *Versioning3Suite) TestChildWorkflowInheritance_UnpinnedParent() { - s.testChildWorkflowInheritance_ExpectNoInherit(false, vbUnpinned) + s.testChildWorkflowInheritanceExpectNoInherit(false, vbUnpinned) } func (s *Versioning3Suite) TestChildWorkflowInheritance_CrossTQ_NoInherit() { // the ExpectNoInherit helper does NOT poll on the child's task queue with the parent's version, // so we DO NOT expect the version to be inherited - s.testChildWorkflowInheritance_ExpectNoInherit(true, vbPinned) + s.testChildWorkflowInheritanceExpectNoInherit(true, vbPinned) } -func (s *Versioning3Suite) testChildWorkflowInheritance_ExpectNoInherit(crossTq bool, parentBehavior enumspb.VersioningBehavior) { +func (s *Versioning3Suite) testChildWorkflowInheritanceExpectNoInherit(crossTq bool, parentBehavior enumspb.VersioningBehavior) { + env := s.setupEnv() + // Child wf of an unpinned parent is always started on the Current Version of its TQ. // For the time being, cross-TQ children do not inherit parents pinned version until that part // is implemented. - tv1 := testvars.New(s).WithBuildIDNumber(1).WithWorkflowIDNumber(1) + tv1 := env.Tv().WithBuildIDNumber(1).WithWorkflowIDNumber(1) tv1Child := tv1.WithWorkflowIDNumber(2) if crossTq { tv1Child = tv1Child.WithTaskQueue("child-tq") @@ -2070,9 +2054,6 @@ func (s *Versioning3Suite) testChildWorkflowInheritance_ExpectNoInherit(crossTq sdkParentBehavior = workflow.VersioningBehaviorAutoUpgrade } - ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second) - defer cancel() - wfStarted := make(chan struct{}, 10) currentChanged := make(chan struct{}, 10) @@ -2095,7 +2076,7 @@ func (s *Versioning3Suite) testChildWorkflowInheritance_ExpectNoInherit(crossTq var val1 string s.NoError(fut1.Get(ctx, &val1)) - s.verifyWorkflowVersioning(s.Assertions, tv1, parentBehavior, tv1.Deployment(), nil, nil) + s.verifyWorkflowVersioning(env, s.Assertions, tv1, parentBehavior, tv1.Deployment(), nil, nil) return val1, nil } @@ -2112,8 +2093,8 @@ func (s *Versioning3Suite) testChildWorkflowInheritance_ExpectNoInherit(crossTq } sdkClient, err := sdkclient.Dial(sdkclient.Options{ - HostPort: s.FrontendGRPCAddress(), - Namespace: s.Namespace().String(), + HostPort: env.FrontendGRPCAddress(), + Namespace: env.Namespace().String(), }) s.NoError(err) @@ -2135,7 +2116,7 @@ func (s *Versioning3Suite) testChildWorkflowInheritance_ExpectNoInherit(crossTq defer w1.Stop() // v1 is current for both parent and child - s.setCurrentDeployment(tv1) + s.setCurrentDeployment(env, tv1) if crossTq { w2xtq := worker.New(sdkClient, tv2Child.TaskQueue().GetName(), worker.Options{ @@ -2165,7 +2146,7 @@ func (s *Versioning3Suite) testChildWorkflowInheritance_ExpectNoInherit(crossTq s.NoError(w2.Start()) defer w2.Stop() - run, err := sdkClient.ExecuteWorkflow(ctx, sdkclient.StartWorkflowOptions{ + run, err := sdkClient.ExecuteWorkflow(env.Context(), sdkclient.StartWorkflowOptions{ ID: tv1.WorkflowID(), TaskQueue: tv1.TaskQueue().GetName(), WorkflowTaskTimeout: 30 * time.Second, @@ -2173,24 +2154,24 @@ func (s *Versioning3Suite) testChildWorkflowInheritance_ExpectNoInherit(crossTq s.NoError(err) // wait for it to start on v1 - s.WaitForChannel(ctx, wfStarted) + env.WaitForChannel(env.Context(), wfStarted) close(wfStarted) // make v2 current for both parent and child and unblock the wf to start the child - s.setCurrentDeployment(tv2) + s.setCurrentDeployment(env, tv2) currentChanged <- struct{}{} var out string - s.NoError(run.Get(ctx, &out)) + s.NoError(run.Get(env.Context(), &out)) s.Equal("v2", out) if parentBehavior == vbPinned { - s.verifyWorkflowVersioning(s.Assertions, tv1, parentBehavior, tv1.Deployment(), nil, nil) + s.verifyWorkflowVersioning(env, s.Assertions, tv1, parentBehavior, tv1.Deployment(), nil, nil) } else { - s.verifyWorkflowVersioning(s.Assertions, tv1, parentBehavior, tv2.Deployment(), nil, nil) + s.verifyWorkflowVersioning(env, s.Assertions, tv1, parentBehavior, tv2.Deployment(), nil, nil) } - s.verifyWorkflowVersioning(s.Assertions, tv2Child, vbPinned, tv2Child.Deployment(), nil, nil) + s.verifyWorkflowVersioning(env, s.Assertions, tv2Child, vbPinned, tv2Child.Deployment(), nil, nil) } func (s *Versioning3Suite) TestPinnedCaN_SameTQ() { @@ -2295,64 +2276,63 @@ func (s *Versioning3Suite) makePinnedOverride(tv *testvars.TestVars) *workflowpb // 5. On WFT: confirm ContinueAsNewSuggested=true, issue ContinueAsNew with AUTO_UPGRADE // 6. The new run should start on v2 (current) and be pinned after WFT completion. func (s *Versioning3Suite) testPinnedCaNUpgradeOnCaN(normalTask, speculativeTask, transientTask, pinnedOverride, enableSendTargetVersionChanged bool) { + var opts []testcore.TestOption if !enableSendTargetVersionChanged { - s.OverrideDynamicConfig(dynamicconfig.EnableSendTargetVersionChanged, false) + opts = append(opts, testcore.WithDynamicConfig(dynamicconfig.EnableSendTargetVersionChanged, false)) } - s.RunTestWithMatchingBehavior(func() { - ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second) - defer cancel() - tv1 := testvars.New(s).WithBuildIDNumber(1) + s.runTestWithMatchingBehavior(func(env *testcore.TestEnv, s *Versioning3Suite) { + tv1 := env.Tv().WithBuildIDNumber(1) tv2 := tv1.WithBuildIDNumber(2) // Set v1 pinned override if testing that var override *workflowpb.VersioningOverride var execution *commonpb.WorkflowExecution if pinnedOverride { override = s.makePinnedOverride(tv1) - execution, _ = s.drainWorkflowTaskAfterSetCurrentWithOverride(tv1, override) + execution, _ = s.drainWorkflowTaskAfterSetCurrentWithOverride(env, tv1, override) } else { - execution, _ = s.drainWorkflowTaskAfterSetCurrent(tv1) + execution, _ = s.drainWorkflowTaskAfterSetCurrent(env, tv1) } // Trigger a normal WFT to make the workflow pinned - s.triggerNormalWFT(ctx, tv1, execution) + s.triggerNormalWFT(env, tv1, execution) // Process the signal WFT and declare behavior to make the workflow versioned pinned on v1 - s.pollWftAndHandle(tv1, false, nil, + s.pollWftAndHandle(env, tv1, false, nil, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) return respondEmptyWft(tv1, false, vbPinned), nil }) // Verify workflow is now pinned and running on v1 - s.verifyWorkflowVersioning(s.Assertions, tv1, vbPinned, tv1.Deployment(), override, nil) - s.verifyVersioningSAs(tv1, vbPinned, enumspb.WORKFLOW_EXECUTION_STATUS_RUNNING, tv1) + s.verifyWorkflowVersioning(env, s.Assertions, tv1, vbPinned, tv1.Deployment(), override, nil) + s.verifyVersioningSAs(env, tv1, vbPinned, enumspb.WORKFLOW_EXECUTION_STATUS_RUNNING, tv1) - s.pollUntilRegistered(ctx, tv2) - s.setCurrentDeployment(tv2) + s.pollUntilRegistered(env, tv2) + s.setCurrentDeployment(env, tv2) // Mode-specific: trigger the WFT var updateResultCh <-chan *workflowservice.UpdateWorkflowExecutionResponse if normalTask { - s.triggerNormalWFT(ctx, tv1, execution) + s.triggerNormalWFT(env, tv1, execution) } else if speculativeTask { - updateResultCh = sendUpdateNoError(s, tv1) + updateResultCh = sendUpdateNoError(env, tv1) } else if transientTask { - s.triggerTransientWFT(ctx, tv1, execution) + s.triggerTransientWFT(env, tv1, execution) } // Process the WFT and issue ContinueAsNew with AUTO_UPGRADE - s.pollWftAndHandle(tv1, false, nil, + s.pollWftAndHandle(env, tv1, false, nil, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) // Mode-specific validations historyEvents := task.History.GetEvents() if speculativeTask { - s.verifySpeculativeTask(execution) + s.verifySpeculativeTask(env, execution) } else if transientTask { s.verifyTransientTask(task) // Get events from server-side history, this includes transient events. - historyEvents = s.GetHistory(s.Namespace().String(), execution) + historyEvents = env.GetHistory(env.Namespace().String(), execution) } wfTaskStartedEvents := make([]*historypb.HistoryEvent, 0) @@ -2368,7 +2348,7 @@ func (s *Versioning3Suite) testPinnedCaNUpgradeOnCaN(normalTask, speculativeTask // The flag is recomputed on every WFT, so both failed attempts and retries will have it if conditions persist. s.Greater(len(wfTaskStartedEvents), 2) // make sure there are at least 3 WFT started events - // In this test, deployment is changed after event 7 (`s.setCurrentDeployment(tv2)`). + // In this test, deployment is changed after event 7 (`s.setCurrentDeployment(env, tv2)`). // So the first 2 WFT started events should NOT have the flag, // and all subsequent events SHOULD have the flag. eventsBeforeDeploymentChange := 2 // Events 3 and 7 @@ -2378,13 +2358,13 @@ func (s *Versioning3Suite) testPinnedCaNUpgradeOnCaN(normalTask, speculativeTask if i < eventsBeforeDeploymentChange { // Events before deployment change should NOT have the flag s.False(attr.GetSuggestContinueAsNew()) - s.Require().Empty(attr.GetSuggestContinueAsNewReasons()) + s.Empty(attr.GetSuggestContinueAsNewReasons()) s.False(attr.GetTargetWorkerDeploymentVersionChanged(), "Event %d should not have flag (before deployment change)", event.GetEventId()) } else { // Events after deployment change SHOULD have the flag (including failed attempts and transient retries) s.False(attr.GetSuggestContinueAsNew()) - s.Require().Empty(attr.GetSuggestContinueAsNewReasons()) + s.Empty(attr.GetSuggestContinueAsNewReasons()) s.True(attr.GetTargetWorkerDeploymentVersionChanged(), "Event %d should have flag (after deployment change)", event.GetEventId()) } @@ -2393,7 +2373,7 @@ func (s *Versioning3Suite) testPinnedCaNUpgradeOnCaN(normalTask, speculativeTask for _, event := range wfTaskStartedEvents { attr := event.GetWorkflowTaskStartedEventAttributes() s.False(attr.GetSuggestContinueAsNew()) - s.Require().Empty(attr.GetSuggestContinueAsNewReasons()) + s.Empty(attr.GetSuggestContinueAsNewReasons()) s.False(attr.GetTargetWorkerDeploymentVersionChanged()) } } @@ -2421,8 +2401,8 @@ func (s *Versioning3Suite) testPinnedCaNUpgradeOnCaN(normalTask, speculativeTask // For speculative WFT, include update accept/complete in the commands if speculativeTask { - resp.Commands = append(s.UpdateAcceptCompleteCommands(tv1), commands...) - resp.Messages = s.UpdateAcceptCompleteMessages(tv1, task.Messages[0]) + resp.Commands = append(s.updateUtils().UpdateAcceptCompleteCommands(tv1), commands...) + resp.Messages = s.updateUtils().UpdateAcceptCompleteMessages(tv1, task.Messages[0]) } return resp, nil @@ -2441,20 +2421,20 @@ func (s *Versioning3Suite) testPinnedCaNUpgradeOnCaN(normalTask, speculativeTask // Start async poller to receive the ContinueAsNew new run wftNewRunDone := make(chan struct{}) - s.pollWftAndHandle(postCaNTV, false, wftNewRunDone, + s.pollWftAndHandle(env, postCaNTV, false, wftNewRunDone, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) return respondCompleteWorkflow(postCaNTV, vbPinned), nil }) // Wait for the new run's first workflow task - s.WaitForChannel(ctx, wftNewRunDone) + env.WaitForChannel(env.Context(), wftNewRunDone) // Verify the new workflow run is on v2 (not v1) because of AUTO_UPGRADE initial behavior // But if testing with Pinned V1 override, it should be on v1 - s.verifyWorkflowVersioning(s.Assertions, postCaNTV, vbPinned, postCaNTV.Deployment(), override, nil) - s.verifyVersioningSAs(postCaNTV, vbPinned, enumspb.WORKFLOW_EXECUTION_STATUS_COMPLETED, postCaNTV) - }) + s.verifyWorkflowVersioning(env, s.Assertions, postCaNTV, vbPinned, postCaNTV.Deployment(), override, nil) + s.verifyVersioningSAs(env, postCaNTV, vbPinned, enumspb.WORKFLOW_EXECUTION_STATUS_COMPLETED, postCaNTV) + }, opts...) } // TestPinnedCaN_UseRampingVersionOnCaN tests that a Pinned workflow's CaN lands on the ramping version. @@ -2493,11 +2473,8 @@ func (s *Versioning3Suite) TestPinnedCaN_UseRampingVersionOnCaN_NoRampingVersion // - noRampingVersion=true: no ramping version exists, falls back to current (v1) // - pinnedOverride=true: override wins, new run stays on v1 func (s *Versioning3Suite) testPinnedCaNUseRampingVersionOnCaN(pinnedOverride, noRampingVersion bool) { - s.RunTestWithMatchingBehavior(func() { - ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second) - defer cancel() - - tv1 := testvars.New(s).WithBuildIDNumber(1) + s.runTestWithMatchingBehavior(func(env *testcore.TestEnv, s *Versioning3Suite) { + tv1 := env.Tv().WithBuildIDNumber(1) tv2 := tv1.WithBuildIDNumber(2) // Set up v1 as current, start workflow @@ -2505,34 +2482,34 @@ func (s *Versioning3Suite) testPinnedCaNUseRampingVersionOnCaN(pinnedOverride, n var execution *commonpb.WorkflowExecution if pinnedOverride { override = s.makePinnedOverride(tv1) - execution, _ = s.drainWorkflowTaskAfterSetCurrentWithOverride(tv1, override) + execution, _ = s.drainWorkflowTaskAfterSetCurrentWithOverride(env, tv1, override) } else { - execution, _ = s.drainWorkflowTaskAfterSetCurrent(tv1) + execution, _ = s.drainWorkflowTaskAfterSetCurrent(env, tv1) } // Trigger a normal WFT to make the workflow pinned on v1 - s.triggerNormalWFT(ctx, tv1, execution) - s.pollWftAndHandle(tv1, false, nil, + s.triggerNormalWFT(env, tv1, execution) + s.pollWftAndHandle(env, tv1, false, nil, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) return respondEmptyWft(tv1, false, vbPinned), nil }) - s.verifyWorkflowVersioning(s.Assertions, tv1, vbPinned, tv1.Deployment(), override, nil) - s.verifyVersioningSAs(tv1, vbPinned, enumspb.WORKFLOW_EXECUTION_STATUS_RUNNING, tv1) + s.verifyWorkflowVersioning(env, s.Assertions, tv1, vbPinned, tv1.Deployment(), override, nil) + s.verifyVersioningSAs(env, tv1, vbPinned, enumspb.WORKFLOW_EXECUTION_STATUS_RUNNING, tv1) if !noRampingVersion { // Register v2 poller before setting it as ramping - s.idlePollWorkflow(ctx, tv2, true, ver3MinPollTime, "should not get any tasks yet") + s.idlePollWorkflow(env, env.Context(), tv2, true, ver3MinPollTime, "should not get any tasks yet") // Set v2 as ramping at 0%: no workflows move via hash, only via UseRampingVersion CaN - s.setRampingDeployment(tv2, 0, false) - s.waitForDeploymentDataPropagation(tv2, versionStatusRamping, false, tqTypeWf) + s.setRampingDeployment(env, tv2, 0, false) + s.waitForDeploymentDataPropagation(env, tv2, versionStatusRamping, false, tqTypeWf) } // Trigger the WFT that will issue the CaN - s.triggerNormalWFT(ctx, tv1, execution) + s.triggerNormalWFT(env, tv1, execution) // Process the WFT and issue ContinueAsNew with USE_RAMPING_VERSION - s.pollWftAndHandle(tv1, false, nil, + s.pollWftAndHandle(env, tv1, false, nil, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) @@ -2540,7 +2517,7 @@ func (s *Versioning3Suite) testPinnedCaNUseRampingVersionOnCaN(pinnedOverride, n if event.GetEventType() == enumspb.EVENT_TYPE_WORKFLOW_TASK_STARTED { attr := event.GetWorkflowTaskStartedEventAttributes() s.False(attr.GetSuggestContinueAsNew()) - s.Require().Empty(attr.GetSuggestContinueAsNewReasons()) + s.Empty(attr.GetSuggestContinueAsNewReasons()) // Setting a ramping version does not change the target for a Pinned workflow: // the target is the pinned version, not current or ramping. s.False(attr.GetTargetWorkerDeploymentVersionChanged(), @@ -2579,15 +2556,15 @@ func (s *Versioning3Suite) testPinnedCaNUseRampingVersionOnCaN(pinnedOverride, n // Receive the new run's first WFT on postCaNTV and complete it pinned wftNewRunDone := make(chan struct{}) - s.pollWftAndHandle(postCaNTV, false, wftNewRunDone, + s.pollWftAndHandle(env, postCaNTV, false, wftNewRunDone, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) return respondCompleteWorkflow(postCaNTV, vbPinned), nil }) - s.WaitForChannel(ctx, wftNewRunDone) + env.WaitForChannel(env.Context(), wftNewRunDone) - s.verifyWorkflowVersioning(s.Assertions, postCaNTV, vbPinned, postCaNTV.Deployment(), override, nil) - s.verifyVersioningSAs(postCaNTV, vbPinned, enumspb.WORKFLOW_EXECUTION_STATUS_COMPLETED, postCaNTV) + s.verifyWorkflowVersioning(env, s.Assertions, postCaNTV, vbPinned, postCaNTV.Deployment(), override, nil) + s.verifyVersioningSAs(env, postCaNTV, vbPinned, enumspb.WORKFLOW_EXECUTION_STATUS_COMPLETED, postCaNTV) }) } @@ -2597,25 +2574,22 @@ func (s *Versioning3Suite) testPinnedCaNUseRampingVersionOnCaN(pinnedOverride, n // Target Version (current), not the ramping version — UseRampingVersion only applies once. // 2. A CaN issued by that run does not inherit UseRampingVersion; its new run uses normal routing. func (s *Versioning3Suite) TestPinnedCaN_UseRampingVersionOnCaN_SubsequentWFTGoesToTargetAndCaNDoesNotInherit() { - s.RunTestWithMatchingBehavior(func() { - ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second) - defer cancel() - - tv1 := testvars.New(s).WithBuildIDNumber(1) + s.runTestWithMatchingBehavior(func(env *testcore.TestEnv, s *Versioning3Suite) { + tv1 := env.Tv().WithBuildIDNumber(1) tv2 := tv1.WithBuildIDNumber(2) // Standard UseRampingVersion CaN setup: pinned v1 workflow CaNs to v2 (ramping at 0%) - execution, _ := s.drainWorkflowTaskAfterSetCurrent(tv1) - s.triggerNormalWFT(ctx, tv1, execution) - s.pollWftAndHandle(tv1, false, nil, + execution, _ := s.drainWorkflowTaskAfterSetCurrent(env, tv1) + s.triggerNormalWFT(env, tv1, execution) + s.pollWftAndHandle(env, tv1, false, nil, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { return respondEmptyWft(tv1, false, vbPinned), nil }) - s.idlePollWorkflow(ctx, tv2, true, ver3MinPollTime, "should not get any tasks yet") - s.setRampingDeployment(tv2, 0, false) - s.waitForDeploymentDataPropagation(tv2, versionStatusRamping, false, tqTypeWf) - s.triggerNormalWFT(ctx, tv1, execution) - s.pollWftAndHandle(tv1, false, nil, + s.idlePollWorkflow(env, env.Context(), tv2, true, ver3MinPollTime, "should not get any tasks yet") + s.setRampingDeployment(env, tv2, 0, false) + s.waitForDeploymentDataPropagation(env, tv2, versionStatusRamping, false, tqTypeWf) + s.triggerNormalWFT(env, tv1, execution) + s.pollWftAndHandle(env, tv1, false, nil, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { return &workflowservice.RespondWorkflowTaskCompletedRequest{ Commands: []*commandpb.Command{{ @@ -2636,17 +2610,17 @@ func (s *Versioning3Suite) TestPinnedCaN_UseRampingVersionOnCaN_SubsequentWFTGoe // First WFT of the UseRampingVersion run lands on v2; worker declares AUTO_UPGRADE. firstWFTDone := make(chan struct{}) - s.pollWftAndHandle(tv2, false, firstWFTDone, + s.pollWftAndHandle(env, tv2, false, firstWFTDone, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) return respondEmptyWft(tv2, false, vbUnpinned), nil }) - s.WaitForChannel(ctx, firstWFTDone) + env.WaitForChannel(env.Context(), firstWFTDone) // Signal the UseRampingVersion run (omit run ID to target the current run). // The resulting WFT should go to v1 (Target/current), not v2 (ramping) — (1). - _, err := s.FrontendClient().SignalWorkflowExecution(ctx, &workflowservice.SignalWorkflowExecutionRequest{ - Namespace: s.Namespace().String(), + _, err := env.FrontendClient().SignalWorkflowExecution(env.Context(), &workflowservice.SignalWorkflowExecutionRequest{ + Namespace: env.Namespace().String(), WorkflowExecution: &commonpb.WorkflowExecution{WorkflowId: tv1.WorkflowID()}, SignalName: tv1.SignalName(), Input: tv1.Any().Payloads(), @@ -2657,7 +2631,7 @@ func (s *Versioning3Suite) TestPinnedCaN_UseRampingVersionOnCaN_SubsequentWFTGoe // Handle the subsequent WFT on v1 and issue a CaN with AUTO_UPGRADE (not USE_RAMPING_VERSION). // This verifies both (1) — v1 receives the WFT — and sets up the CaN inheritance check (2). secondWFTDone := make(chan struct{}) - s.pollWftAndHandle(tv1, false, secondWFTDone, + s.pollWftAndHandle(env, tv1, false, secondWFTDone, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) return &workflowservice.RespondWorkflowTaskCompletedRequest{ @@ -2676,20 +2650,20 @@ func (s *Versioning3Suite) TestPinnedCaN_UseRampingVersionOnCaN_SubsequentWFTGoe DeploymentOptions: tv2.WorkerDeploymentOptions(true), }, nil }) - s.WaitForChannel(ctx, secondWFTDone) + env.WaitForChannel(env.Context(), secondWFTDone) // The CaN run's first WFT must land on v1 (current), NOT v2 (ramping) — (2). // UseRampingVersion is not inherited by a CaN issued from the resulting workflow. thirdRunDone := make(chan struct{}) - s.pollWftAndHandle(tv1, false, thirdRunDone, + s.pollWftAndHandle(env, tv1, false, thirdRunDone, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) return respondCompleteWorkflow(tv1, vbUnpinned), nil }) - s.WaitForChannel(ctx, thirdRunDone) + env.WaitForChannel(env.Context(), thirdRunDone) - s.verifyWorkflowVersioning(s.Assertions, tv1, vbUnpinned, tv1.Deployment(), nil, nil) - s.verifyVersioningSAs(tv1, vbUnpinned, enumspb.WORKFLOW_EXECUTION_STATUS_COMPLETED, tv1) + s.verifyWorkflowVersioning(env, s.Assertions, tv1, vbUnpinned, tv1.Deployment(), nil, nil) + s.verifyVersioningSAs(env, tv1, vbUnpinned, enumspb.WORKFLOW_EXECUTION_STATUS_COMPLETED, tv1) }) } @@ -2697,25 +2671,22 @@ func (s *Versioning3Suite) TestPinnedCaN_UseRampingVersionOnCaN_SubsequentWFTGoe // UseRampingVersion run also lands on the ramping version for its first WFT. // This verifies that ContinueAsNewInitialVersioningBehavior is propagated through retries per spec. func (s *Versioning3Suite) TestPinnedCaN_UseRampingVersionOnCaN_RetryInheritsInitialBehavior() { - s.RunTestWithMatchingBehavior(func() { - ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second) - defer cancel() - - tv1 := testvars.New(s).WithBuildIDNumber(1) + s.runTestWithMatchingBehavior(func(env *testcore.TestEnv, s *Versioning3Suite) { + tv1 := env.Tv().WithBuildIDNumber(1) tv2 := tv1.WithBuildIDNumber(2) // Standard UseRampingVersion CaN setup, with a retry policy so the run can retry on failure. - execution, _ := s.drainWorkflowTaskAfterSetCurrent(tv1) - s.triggerNormalWFT(ctx, tv1, execution) - s.pollWftAndHandle(tv1, false, nil, + execution, _ := s.drainWorkflowTaskAfterSetCurrent(env, tv1) + s.triggerNormalWFT(env, tv1, execution) + s.pollWftAndHandle(env, tv1, false, nil, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { return respondEmptyWft(tv1, false, vbPinned), nil }) - s.idlePollWorkflow(ctx, tv2, true, ver3MinPollTime, "should not get any tasks yet") - s.setRampingDeployment(tv2, 0, false) - s.waitForDeploymentDataPropagation(tv2, versionStatusRamping, false, tqTypeWf) - s.triggerNormalWFT(ctx, tv1, execution) - s.pollWftAndHandle(tv1, false, nil, + s.idlePollWorkflow(env, env.Context(), tv2, true, ver3MinPollTime, "should not get any tasks yet") + s.setRampingDeployment(env, tv2, 0, false) + s.waitForDeploymentDataPropagation(env, tv2, versionStatusRamping, false, tqTypeWf) + s.triggerNormalWFT(env, tv1, execution) + s.pollWftAndHandle(env, tv1, false, nil, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { return &workflowservice.RespondWorkflowTaskCompletedRequest{ Commands: []*commandpb.Command{{ @@ -2738,7 +2709,7 @@ func (s *Versioning3Suite) TestPinnedCaN_UseRampingVersionOnCaN_RetryInheritsIni // First attempt: first WFT of the UseRampingVersion run lands on v2. // Declares AUTO_UPGRADE and then immediately fails (retryable). firstAttemptDone := make(chan struct{}) - s.pollWftAndHandle(tv2, false, firstAttemptDone, + s.pollWftAndHandle(env, tv2, false, firstAttemptDone, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) return &workflowservice.RespondWorkflowTaskCompletedRequest{ @@ -2754,20 +2725,20 @@ func (s *Versioning3Suite) TestPinnedCaN_UseRampingVersionOnCaN_RetryInheritsIni DeploymentOptions: tv2.WorkerDeploymentOptions(true), }, nil }) - s.WaitForChannel(ctx, firstAttemptDone) + env.WaitForChannel(env.Context(), firstAttemptDone) // Retry run: ContinueAsNewInitialVersioningBehavior=USE_RAMPING_VERSION is propagated, // so the retry's first WFT must also land on v2 (ramping), not v1 (current). retryRunDone := make(chan struct{}) - s.pollWftAndHandle(tv2, false, retryRunDone, + s.pollWftAndHandle(env, tv2, false, retryRunDone, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) return respondCompleteWorkflow(tv2, vbUnpinned), nil }) - s.WaitForChannel(ctx, retryRunDone) + env.WaitForChannel(env.Context(), retryRunDone) - s.verifyWorkflowVersioning(s.Assertions, tv2, vbUnpinned, tv2.Deployment(), nil, nil) - s.verifyVersioningSAs(tv2, vbUnpinned, enumspb.WORKFLOW_EXECUTION_STATUS_COMPLETED, tv2) + s.verifyWorkflowVersioning(env, s.Assertions, tv2, vbUnpinned, tv2.Deployment(), nil, nil) + s.verifyVersioningSAs(env, tv2, vbUnpinned, enumspb.WORKFLOW_EXECUTION_STATUS_COMPLETED, tv2) }) } @@ -2776,26 +2747,23 @@ func (s *Versioning3Suite) TestPinnedCaN_UseRampingVersionOnCaN_RetryInheritsIni // as its parent based on InheritedAutoUpgradeInfo, and therefore starts on v2 also, but we confirm that // it does not inherit UseRampingVersion. func (s *Versioning3Suite) TestPinnedCaN_UseRampingVersionOnCaN_ChildDoesNotInheritUseRampingVersion() { - s.RunTestWithMatchingBehavior(func() { - ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second) - defer cancel() - - tv1 := testvars.New(s).WithBuildIDNumber(1) + s.runTestWithMatchingBehavior(func(env *testcore.TestEnv, s *Versioning3Suite) { + tv1 := env.Tv().WithBuildIDNumber(1) tv2 := tv1.WithBuildIDNumber(2) childID := tv1.WorkflowID() + "-child" // Standard UseRampingVersion CaN setup. - execution, _ := s.drainWorkflowTaskAfterSetCurrent(tv1) - s.triggerNormalWFT(ctx, tv1, execution) - s.pollWftAndHandle(tv1, false, nil, + execution, _ := s.drainWorkflowTaskAfterSetCurrent(env, tv1) + s.triggerNormalWFT(env, tv1, execution) + s.pollWftAndHandle(env, tv1, false, nil, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { return respondEmptyWft(tv1, false, vbPinned), nil }) - s.idlePollWorkflow(ctx, tv2, true, ver3MinPollTime, "should not get any tasks yet") - s.setRampingDeployment(tv2, 0, false) - s.waitForDeploymentDataPropagation(tv2, versionStatusRamping, false, tqTypeWf) - s.triggerNormalWFT(ctx, tv1, execution) - s.pollWftAndHandle(tv1, false, nil, + s.idlePollWorkflow(env, env.Context(), tv2, true, ver3MinPollTime, "should not get any tasks yet") + s.setRampingDeployment(env, tv2, 0, false) + s.waitForDeploymentDataPropagation(env, tv2, versionStatusRamping, false, tqTypeWf) + s.triggerNormalWFT(env, tv1, execution) + s.pollWftAndHandle(env, tv1, false, nil, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { return &workflowservice.RespondWorkflowTaskCompletedRequest{ Commands: []*commandpb.Command{{ @@ -2816,7 +2784,7 @@ func (s *Versioning3Suite) TestPinnedCaN_UseRampingVersionOnCaN_ChildDoesNotInhe // First WFT of the UseRampingVersion run on v2: declare AUTO_UPGRADE and start a child workflow. parentWFT1Done := make(chan struct{}) - s.pollWftAndHandle(tv2, false, parentWFT1Done, + s.pollWftAndHandle(env, tv2, false, parentWFT1Done, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) return &workflowservice.RespondWorkflowTaskCompletedRequest{ @@ -2835,47 +2803,26 @@ func (s *Versioning3Suite) TestPinnedCaN_UseRampingVersionOnCaN_ChildDoesNotInhe DeploymentOptions: tv2.WorkerDeploymentOptions(true), }, nil }) - s.WaitForChannel(ctx, parentWFT1Done) + env.WaitForChannel(env.Context(), parentWFT1Done) - // Child's first WFT lands on v2 because of normal InheritedAutoUpgradeInfo propagation. // Child workflows do NOT inherit UseRampingVersion. childWFTDone := make(chan struct{}) - s.pollWftAndHandle(tv1, false, childWFTDone, - func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { - s.NotNil(task) - return respondCompleteWorkflow(tv1, vbUnpinned), nil - }) - s.WaitForChannel(ctx, childWFTDone) - - // Parent receives a WFT after the child completes; complete the parent too. - parentWFT2Done := make(chan struct{}) - s.pollWftAndHandle(tv1, false, parentWFT2Done, + s.pollWftAndHandle(env, tv1, false, childWFTDone, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) return respondCompleteWorkflow(tv1, vbUnpinned), nil }) - s.WaitForChannel(ctx, parentWFT2Done) + env.WaitForChannel(env.Context(), childWFTDone) - // Verify the child ended on v2 because of normal InheritedAutoUpgradeInfo propagation, - // not because of UseRampingVersion. + // Verify the child did not inherit UseRampingVersion. s.EventuallyWithT(func(t *assert.CollectT) { - resp, err := s.FrontendClient().DescribeWorkflowExecution(ctx, &workflowservice.DescribeWorkflowExecutionRequest{ - Namespace: s.Namespace().String(), - Execution: &commonpb.WorkflowExecution{WorkflowId: childID}, - }) - s.NoError(err) - s.Equal(tv2.Deployment().GetBuildId(), - resp.GetWorkflowExecutionInfo().GetVersioningInfo().GetDeploymentVersion().GetBuildId()) - s.Equal(tv2.Deployment().GetSeriesName(), - resp.GetWorkflowExecutionInfo().GetVersioningInfo().GetDeploymentVersion().GetDeploymentName()) - - resp2, err := s.FrontendClient().GetWorkflowExecutionHistory(ctx, &workflowservice.GetWorkflowExecutionHistoryRequest{ - Namespace: s.Namespace().String(), + resp, err := env.FrontendClient().GetWorkflowExecutionHistory(env.Context(), &workflowservice.GetWorkflowExecutionHistoryRequest{ + Namespace: env.Namespace().String(), Execution: &commonpb.WorkflowExecution{WorkflowId: childID}, }) s.NoError(err) foundStartedEvent := false - for _, e := range resp2.GetHistory().GetEvents() { + for _, e := range resp.GetHistory().GetEvents() { if e.GetEventType() == enumspb.EVENT_TYPE_WORKFLOW_EXECUTION_STARTED { foundStartedEvent = true s.Equal(enumspb.CONTINUE_AS_NEW_VERSIONING_BEHAVIOR_UNSPECIFIED, @@ -2888,9 +2835,9 @@ func (s *Versioning3Suite) TestPinnedCaN_UseRampingVersionOnCaN_ChildDoesNotInhe } // Signal to trigger a normal WFT -func (s *Versioning3Suite) triggerNormalWFT(ctx context.Context, tv *testvars.TestVars, execution *commonpb.WorkflowExecution) { - _, err := s.FrontendClient().SignalWorkflowExecution(ctx, &workflowservice.SignalWorkflowExecutionRequest{ - Namespace: s.Namespace().String(), +func (s *Versioning3Suite) triggerNormalWFT(env *testcore.TestEnv, tv *testvars.TestVars, execution *commonpb.WorkflowExecution) { + _, err := env.FrontendClient().SignalWorkflowExecution(env.Context(), &workflowservice.SignalWorkflowExecutionRequest{ + Namespace: env.Namespace().String(), WorkflowExecution: execution, SignalName: tv.SignalName(), Input: tv.Any().Payloads(), @@ -2900,12 +2847,12 @@ func (s *Versioning3Suite) triggerNormalWFT(ctx context.Context, tv *testvars.Te } // Trigger a normal task and then fail the task twice to trigger a transient WFT -func (s *Versioning3Suite) triggerTransientWFT(ctx context.Context, tv *testvars.TestVars, execution *commonpb.WorkflowExecution) { - s.triggerNormalWFT(ctx, tv, execution) +func (s *Versioning3Suite) triggerTransientWFT(env *testcore.TestEnv, tv *testvars.TestVars, execution *commonpb.WorkflowExecution) { + s.triggerNormalWFT(env, tv, execution) // Poll and FAIL the WFT to create a transient WFT situation - pollResp, err := s.FrontendClient().PollWorkflowTaskQueue(ctx, &workflowservice.PollWorkflowTaskQueueRequest{ - Namespace: s.Namespace().String(), + pollResp, err := env.FrontendClient().PollWorkflowTaskQueue(env.Context(), &workflowservice.PollWorkflowTaskQueueRequest{ + Namespace: env.Namespace().String(), TaskQueue: tv.TaskQueue(), Identity: tv.WorkerIdentity(), DeploymentOptions: tv.WorkerDeploymentOptions(true), @@ -2915,8 +2862,8 @@ func (s *Versioning3Suite) triggerTransientWFT(ctx context.Context, tv *testvars s.NotEmpty(pollResp.TaskToken) // Fail the workflow task - this will cause a transient WFT to be scheduled - _, err = s.FrontendClient().RespondWorkflowTaskFailed(ctx, &workflowservice.RespondWorkflowTaskFailedRequest{ - Namespace: s.Namespace().String(), + _, err = env.FrontendClient().RespondWorkflowTaskFailed(env.Context(), &workflowservice.RespondWorkflowTaskFailedRequest{ + Namespace: env.Namespace().String(), TaskToken: pollResp.TaskToken, Cause: enumspb.WORKFLOW_TASK_FAILED_CAUSE_WORKFLOW_WORKER_UNHANDLED_FAILURE, Identity: tv.WorkerIdentity(), @@ -2925,8 +2872,8 @@ func (s *Versioning3Suite) triggerTransientWFT(ctx context.Context, tv *testvars } // Verify this is a speculative task - events not yet in persisted history -func (s *Versioning3Suite) verifySpeculativeTask(execution *commonpb.WorkflowExecution) { - events := s.GetHistory(s.Namespace().String(), execution) +func (s *Versioning3Suite) verifySpeculativeTask(env *testcore.TestEnv, execution *commonpb.WorkflowExecution) { + events := env.GetHistory(env.Namespace().String(), execution) s.EqualHistoryEvents(` 1 WorkflowExecutionStarted 2 WorkflowTaskScheduled @@ -2949,7 +2896,7 @@ func (s *Versioning3Suite) verifyTransientTask(task *workflowservice.PollWorkflo lastScheduledEvent = event } } - s.Require().NotNil(lastScheduledEvent) + s.NotNil(lastScheduledEvent) s.Equal(int32(2), lastScheduledEvent.GetWorkflowTaskScheduledEventAttributes().GetAttempt()) } @@ -2966,23 +2913,20 @@ func (s *Versioning3Suite) verifyTransientTask(task *workflowservice.PollWorkflo // - Issue ContinueAsNew with InitialVersioningBehavior AUTO_UPGRADE // -> first task should be AUTO_UPGRADE (same as it would be without InitialVersioningBehavior) func (s *Versioning3Suite) TestAutoUpgradeCaN_UpgradeOnCaN() { - s.RunTestWithMatchingBehavior(func() { - ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second) - defer cancel() - - tv1 := testvars.New(s).WithBuildIDNumber(1) + s.runTestWithMatchingBehavior(func(env *testcore.TestEnv, s *Versioning3Suite) { + tv1 := env.Tv().WithBuildIDNumber(1) tv2 := tv1.WithBuildIDNumber(2) - execution, _ := s.drainWorkflowTaskAfterSetCurrent(tv1) + execution, _ := s.drainWorkflowTaskAfterSetCurrent(env, tv1) // Verify workflow is now versioned (AutoUpgrade) and running on v1 - s.verifyWorkflowVersioning(s.Assertions, tv1, vbUnpinned, tv1.Deployment(), nil, nil) + s.verifyWorkflowVersioning(env, s.Assertions, tv1, vbUnpinned, tv1.Deployment(), nil, nil) - s.pollUntilRegistered(ctx, tv2) - s.setCurrentDeployment(tv2) + s.pollUntilRegistered(env, tv2) + s.setCurrentDeployment(env, tv2) // Signal the workflow again to trigger the WFT with ContinueAsNewSuggested=true and reasons=[NewTargetVersion] - _, err := s.FrontendClient().SignalWorkflowExecution(ctx, &workflowservice.SignalWorkflowExecutionRequest{ - Namespace: s.Namespace().String(), + _, err := env.FrontendClient().SignalWorkflowExecution(env.Context(), &workflowservice.SignalWorkflowExecutionRequest{ + Namespace: env.Namespace().String(), WorkflowExecution: execution, SignalName: tv1.SignalName(), Input: tv1.Any().Payloads(), @@ -2991,7 +2935,7 @@ func (s *Versioning3Suite) TestAutoUpgradeCaN_UpgradeOnCaN() { s.NoError(err) // Process the signal WFT and issue ContinueAsNew with AUTO_UPGRADE initial behavior - s.pollWftAndHandle(tv2, false, nil, + s.pollWftAndHandle(env, tv2, false, nil, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) @@ -3001,12 +2945,12 @@ func (s *Versioning3Suite) TestAutoUpgradeCaN_UpgradeOnCaN() { wfTaskStartedEvents = append(wfTaskStartedEvents, event) } } - s.Require().Len(wfTaskStartedEvents, 2) // make sure we are actually verifying non-zero # of events + s.Len(wfTaskStartedEvents, 2) // make sure we are actually verifying non-zero # of events for _, event := range wfTaskStartedEvents { attr := event.GetWorkflowTaskStartedEventAttributes() s.False(attr.GetSuggestContinueAsNew()) - s.Require().Empty(attr.GetSuggestContinueAsNewReasons()) + s.Empty(attr.GetSuggestContinueAsNewReasons()) } // For AutoUpgrade, I want to test that once the workflow has transitioned to v2, it doesn't get the CaN suggestion anymore. @@ -3014,15 +2958,15 @@ func (s *Versioning3Suite) TestAutoUpgradeCaN_UpgradeOnCaN() { }) // Signal the workflow again to trigger another WFT with ContinueAsNewSuggested=false and reasons=[] - _, err = s.FrontendClient().SignalWorkflowExecution(ctx, &workflowservice.SignalWorkflowExecutionRequest{ - Namespace: s.Namespace().String(), + _, err = env.FrontendClient().SignalWorkflowExecution(env.Context(), &workflowservice.SignalWorkflowExecutionRequest{ + Namespace: env.Namespace().String(), WorkflowExecution: execution, SignalName: tv1.SignalName(), Input: tv1.Any().Payloads(), Identity: tv1.WorkerIdentity(), }) s.NoError(err) - s.pollWftAndHandle(tv2, false, nil, + s.pollWftAndHandle(env, tv2, false, nil, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) @@ -3037,7 +2981,7 @@ func (s *Versioning3Suite) TestAutoUpgradeCaN_UpgradeOnCaN() { for _, event := range wfTaskStartedEvents { attr := event.GetWorkflowTaskStartedEventAttributes() s.False(attr.GetSuggestContinueAsNew()) - s.Require().Empty(attr.GetSuggestContinueAsNewReasons()) + s.Empty(attr.GetSuggestContinueAsNewReasons()) } return &workflowservice.RespondWorkflowTaskCompletedRequest{ @@ -3062,7 +3006,7 @@ func (s *Versioning3Suite) TestAutoUpgradeCaN_UpgradeOnCaN() { // Start async poller for v2 to receive the ContinueAsNew new run wftNewRunDone := make(chan struct{}) - s.pollWftAndHandle(tv2, false, wftNewRunDone, + s.pollWftAndHandle(env, tv2, false, wftNewRunDone, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) // The new run should be on v2 because InitialVersioningBehavior was AUTO_UPGRADE @@ -3070,17 +3014,19 @@ func (s *Versioning3Suite) TestAutoUpgradeCaN_UpgradeOnCaN() { }) // Wait for the new run's first workflow task on v2 - s.WaitForChannel(ctx, wftNewRunDone) + env.WaitForChannel(env.Context(), wftNewRunDone) // Verify the new workflow run is on v2 (not v1) because of AUTO_UPGRADE initial behavior // The new workflow has the worker-sent AutoUpgrade behavior from the first WFT completion after CaN. - s.verifyWorkflowVersioning(s.Assertions, tv2, vbUnpinned, tv2.Deployment(), nil, nil) + s.verifyWorkflowVersioning(env, s.Assertions, tv2, vbUnpinned, tv2.Deployment(), nil, nil) }) } func (s *Versioning3Suite) testCan(crossTq bool, behavior enumspb.VersioningBehavior, upgradeOnCaN bool, expectPinnedInherit bool, pinnedOverride bool) { + env := s.setupEnv() + // CaN inherits version if pinned and if new task queue is in pinned version, goes to current version if unpinned. - tv1 := testvars.New(s).WithBuildIDNumber(1).WithWorkflowIDNumber(1) + tv1 := env.Tv().WithBuildIDNumber(1).WithWorkflowIDNumber(1) tv2 := tv1.WithBuildIDNumber(2) sdkBehavior := workflow.VersioningBehaviorAutoUpgrade if behavior == vbPinned { @@ -3093,9 +3039,6 @@ func (s *Versioning3Suite) testCan(crossTq bool, behavior enumspb.VersioningBeha override = s.makePinnedOverride(tv1) } - ctx, cancel := context.WithTimeout(context.Background(), 600*time.Second) - defer cancel() - wfStarted := make(chan struct{}, 10) currentChanged := make(chan struct{}, 10) @@ -3111,17 +3054,17 @@ func (s *Versioning3Suite) testCan(crossTq bool, behavior enumspb.VersioningBeha // newCtx = workflow.WithInitialVersioningBehavior(newCtx, temporal.ContinueAsNewVersioningBehaviorAutoUpgrade) // } if pinnedOverride { - s.verifyWorkflowVersioning(s.Assertions, tv1, vbUnspecified, nil, override, nil) + s.verifyWorkflowVersioning(env, s.Assertions, tv1, vbUnspecified, nil, override, nil) } else { - s.verifyWorkflowVersioning(s.Assertions, tv1, vbUnspecified, nil, override, tv1.DeploymentVersionTransition()) + s.verifyWorkflowVersioning(env, s.Assertions, tv1, vbUnspecified, nil, override, tv1.DeploymentVersionTransition()) } wfStarted <- struct{}{} // wait for current version to change. <-currentChanged return "", workflow.NewContinueAsNewError(newCtx, "wf", attempt+1) case 1: - s.verifyWorkflowVersioning(s.Assertions, tv1, vbPinned, tv1.Deployment(), override, nil) - s.verifyVersioningSAs(tv1, vbPinned, enumspb.WORKFLOW_EXECUTION_STATUS_RUNNING) + s.verifyWorkflowVersioning(env, s.Assertions, tv1, vbPinned, tv1.Deployment(), override, nil) + s.verifyVersioningSAs(env, tv1, vbPinned, enumspb.WORKFLOW_EXECUTION_STATUS_RUNNING) return "v1", nil } s.FailNow("workflow should not get to this point") @@ -3129,20 +3072,20 @@ func (s *Versioning3Suite) testCan(crossTq bool, behavior enumspb.VersioningBeha } wf2 := func(ctx workflow.Context, attempt int) (string, error) { - if behavior == vbUnpinned && s.deploymentWorkflowVersion >= workerdeployment.AsyncSetCurrentAndRamping { + if behavior == vbUnpinned && versioning3DeploymentWorkflowVersion >= workerdeployment.AsyncSetCurrentAndRamping { // Unpinned CaN should inherit parent deployment version and behaviour - s.verifyWorkflowVersioning(s.Assertions, tv2, vbUnpinned, tv1.Deployment(), override, tv2.DeploymentVersionTransition()) - s.verifyVersioningSAs(tv1, vbUnpinned, enumspb.WORKFLOW_EXECUTION_STATUS_RUNNING) + s.verifyWorkflowVersioning(env, s.Assertions, tv2, vbUnpinned, tv1.Deployment(), override, tv2.DeploymentVersionTransition()) + s.verifyVersioningSAs(env, tv1, vbUnpinned, enumspb.WORKFLOW_EXECUTION_STATUS_RUNNING) } else { - s.verifyWorkflowVersioning(s.Assertions, tv2, vbUnspecified, nil, override, tv2.DeploymentVersionTransition()) - s.verifyVersioningSAs(tv2, vbUnspecified, enumspb.WORKFLOW_EXECUTION_STATUS_RUNNING) + s.verifyWorkflowVersioning(env, s.Assertions, tv2, vbUnspecified, nil, override, tv2.DeploymentVersionTransition()) + s.verifyVersioningSAs(env, tv2, vbUnspecified, enumspb.WORKFLOW_EXECUTION_STATUS_RUNNING) } return "v2", nil } sdkClient, err := sdkclient.Dial(sdkclient.Options{ - HostPort: s.FrontendGRPCAddress(), - Namespace: s.Namespace().String(), + HostPort: env.FrontendGRPCAddress(), + Namespace: env.Namespace().String(), }) s.NoError(err) @@ -3199,7 +3142,7 @@ func (s *Versioning3Suite) testCan(crossTq bool, behavior enumspb.VersioningBeha defer w2.Stop() // make v1 current - s.setCurrentDeployment(tv1) + s.setCurrentDeployment(env, tv1) startOpts := sdkclient.StartWorkflowOptions{ ID: tv1.WorkflowID(), @@ -3215,32 +3158,31 @@ func (s *Versioning3Suite) testCan(crossTq bool, behavior enumspb.VersioningBeha } } - run, err := sdkClient.ExecuteWorkflow(ctx, startOpts, "wf") + run, err := sdkClient.ExecuteWorkflow(env.Context(), startOpts, "wf") s.NoError(err) // wait for it to start on v1 - s.WaitForChannel(ctx, wfStarted) + env.WaitForChannel(env.Context(), wfStarted) close(wfStarted) // force panic if replayed // make v2 current - s.setCurrentDeployment(tv2) + s.setCurrentDeployment(env, tv2) currentChanged <- struct{}{} var out string - s.NoError(run.Get(ctx, &out)) + s.NoError(run.Get(env.Context(), &out)) if expectPinnedInherit { s.Equal("v1", out) - s.verifyVersioningSAs(tv1, vbPinned, enumspb.WORKFLOW_EXECUTION_STATUS_COMPLETED, tv1) + s.verifyVersioningSAs(env, tv1, vbPinned, enumspb.WORKFLOW_EXECUTION_STATUS_COMPLETED, tv1) } else { s.Equal("v2", out) - s.verifyVersioningSAs(tv2, behavior, enumspb.WORKFLOW_EXECUTION_STATUS_COMPLETED, tv2) + s.verifyVersioningSAs(env, tv2, behavior, enumspb.WORKFLOW_EXECUTION_STATUS_COMPLETED, tv2) } } func (s *Versioning3Suite) TestDescribeTaskQueueVersioningInfo() { - tv := testvars.New(s) - ctx, cancel := context.WithTimeout(context.Background(), 50*time.Second) - defer cancel() + env := s.setupEnv() + tv := env.Tv() t1 := time.Now() t2 := t1.Add(time.Second) @@ -3257,11 +3199,11 @@ func (s *Versioning3Suite) TestDescribeTaskQueueVersioningInfo() { RevisionNumber: revisionNumber, } - s.syncTaskQueueDeploymentDataWithRoutingConfig(tv, newRoutingConfig, map[string]*deploymentspb.WorkerDeploymentVersionData{tv.DeploymentVersion().GetBuildId(): { + s.syncTaskQueueDeploymentDataWithRoutingConfig(env, tv, newRoutingConfig, map[string]*deploymentspb.WorkerDeploymentVersionData{tv.DeploymentVersion().GetBuildId(): { Status: enumspb.WORKER_DEPLOYMENT_VERSION_STATUS_CURRENT, }}, []string{}, tqTypeWf) - wfInfo, err := s.FrontendClient().DescribeTaskQueue(ctx, &workflowservice.DescribeTaskQueueRequest{ - Namespace: s.Namespace().String(), + wfInfo, err := env.FrontendClient().DescribeTaskQueue(env.Context(), &workflowservice.DescribeTaskQueueRequest{ + Namespace: env.Namespace().String(), TaskQueue: tv.TaskQueue(), TaskQueueType: tqTypeWf, }) @@ -3281,12 +3223,12 @@ func (s *Versioning3Suite) TestDescribeTaskQueueVersioningInfo() { CurrentVersionChangedTime: timestamp.TimePtr(t1), RevisionNumber: revisionNumber, } - s.syncTaskQueueDeploymentDataWithRoutingConfig(tv, newRoutingConfig, map[string]*deploymentspb.WorkerDeploymentVersionData{tv.DeploymentVersion().GetBuildId(): { + s.syncTaskQueueDeploymentDataWithRoutingConfig(env, tv, newRoutingConfig, map[string]*deploymentspb.WorkerDeploymentVersionData{tv.DeploymentVersion().GetBuildId(): { Status: enumspb.WORKER_DEPLOYMENT_VERSION_STATUS_CURRENT, }}, []string{}, tqTypeAct) - actInfo, err := s.FrontendClient().DescribeTaskQueue(ctx, &workflowservice.DescribeTaskQueueRequest{ - Namespace: s.Namespace().String(), + actInfo, err := env.FrontendClient().DescribeTaskQueue(env.Context(), &workflowservice.DescribeTaskQueueRequest{ + Namespace: env.Namespace().String(), TaskQueue: tv.TaskQueue(), TaskQueueType: tqTypeAct, }) @@ -3307,13 +3249,13 @@ func (s *Versioning3Suite) TestDescribeTaskQueueVersioningInfo() { RampingVersionPercentageChangedTime: timestamp.TimePtr(t2), RevisionNumber: 2, } - s.syncTaskQueueDeploymentDataWithRoutingConfig(tv, newRoutingConfig, map[string]*deploymentspb.WorkerDeploymentVersionData{tv.DeploymentVersion().GetBuildId(): { + s.syncTaskQueueDeploymentDataWithRoutingConfig(env, tv, newRoutingConfig, map[string]*deploymentspb.WorkerDeploymentVersionData{tv.DeploymentVersion().GetBuildId(): { Status: enumspb.WORKER_DEPLOYMENT_VERSION_STATUS_CURRENT, }}, []string{}, tqTypeAct) - s.waitForDeploymentDataPropagation(tv, versionStatusNil, true, tqTypeAct) + s.waitForDeploymentDataPropagation(env, tv, versionStatusNil, true, tqTypeAct) - actInfo, err = s.FrontendClient().DescribeTaskQueue(ctx, &workflowservice.DescribeTaskQueueRequest{ - Namespace: s.Namespace().String(), + actInfo, err = env.FrontendClient().DescribeTaskQueue(env.Context(), &workflowservice.DescribeTaskQueueRequest{ + Namespace: env.Namespace().String(), TaskQueue: tv.TaskQueue(), TaskQueueType: tqTypeAct, }) @@ -3329,11 +3271,12 @@ func (s *Versioning3Suite) TestDescribeTaskQueueVersioningInfo() { } func (s *Versioning3Suite) TestSyncDeploymentUserDataWithRoutingConfig_Update() { - tv := testvars.New(s) + env := s.setupEnv() + tv := env.Tv() - data := s.getTaskQueueDeploymentData(tv, tqTypeAct) + data := s.getTaskQueueDeploymentData(env, tv, tqTypeAct) s.Nil(data) - data = s.getTaskQueueDeploymentData(tv, tqTypeWf) + data = s.getTaskQueueDeploymentData(env, tv, tqTypeWf) s.Nil(data) t1 := time.Now() @@ -3345,11 +3288,11 @@ func (s *Versioning3Suite) TestSyncDeploymentUserDataWithRoutingConfig_Update() CurrentVersionChangedTime: timestamp.TimePtr(t1), RevisionNumber: revisionNumber, } - s.syncTaskQueueDeploymentDataWithRoutingConfig(tv1, routingConfig, map[string]*deploymentspb.WorkerDeploymentVersionData{tv1.DeploymentVersion().GetBuildId(): { + s.syncTaskQueueDeploymentDataWithRoutingConfig(env, tv1, routingConfig, map[string]*deploymentspb.WorkerDeploymentVersionData{tv1.DeploymentVersion().GetBuildId(): { Status: enumspb.WORKER_DEPLOYMENT_VERSION_STATUS_CURRENT, }}, []string{}, tqTypeAct) - data = s.getTaskQueueDeploymentData(tv, tqTypeAct) + data = s.getTaskQueueDeploymentData(env, tv, tqTypeAct) s.ProtoEqual(&persistencespb.DeploymentData{DeploymentsData: map[string]*persistencespb.WorkerDeploymentData{ tv1.DeploymentVersion().GetDeploymentName(): { RoutingConfig: routingConfig, @@ -3358,7 +3301,7 @@ func (s *Versioning3Suite) TestSyncDeploymentUserDataWithRoutingConfig_Update() }}, }, }}, data) - data = s.getTaskQueueDeploymentData(tv, tqTypeWf) + data = s.getTaskQueueDeploymentData(env, tv, tqTypeWf) s.Nil(data) // Changing things with an older revision number should not have effect. @@ -3367,10 +3310,10 @@ func (s *Versioning3Suite) TestSyncDeploymentUserDataWithRoutingConfig_Update() CurrentVersionChangedTime: timestamp.TimePtr(t1), RevisionNumber: revisionNumber - 1, } - s.syncTaskQueueDeploymentDataWithRoutingConfig(tv1, invalidRoutingConfigWithOlderRevisionNumber, map[string]*deploymentspb.WorkerDeploymentVersionData{tv1.DeploymentVersion().GetBuildId(): { + s.syncTaskQueueDeploymentDataWithRoutingConfig(env, tv1, invalidRoutingConfigWithOlderRevisionNumber, map[string]*deploymentspb.WorkerDeploymentVersionData{tv1.DeploymentVersion().GetBuildId(): { Status: enumspb.WORKER_DEPLOYMENT_VERSION_STATUS_CURRENT, }}, []string{}, tqTypeAct) - data = s.getTaskQueueDeploymentData(tv, tqTypeAct) + data = s.getTaskQueueDeploymentData(env, tv, tqTypeAct) s.ProtoEqual(&persistencespb.DeploymentData{DeploymentsData: map[string]*persistencespb.WorkerDeploymentData{ tv1.DeploymentVersion().GetDeploymentName(): { RoutingConfig: routingConfig, @@ -3387,10 +3330,10 @@ func (s *Versioning3Suite) TestSyncDeploymentUserDataWithRoutingConfig_Update() CurrentVersionChangedTime: timestamp.TimePtr(t1), RevisionNumber: revisionNumber, } - s.syncTaskQueueDeploymentDataWithRoutingConfig(tv1, routingConfig, map[string]*deploymentspb.WorkerDeploymentVersionData{tv1.DeploymentVersion().GetBuildId(): { + s.syncTaskQueueDeploymentDataWithRoutingConfig(env, tv1, routingConfig, map[string]*deploymentspb.WorkerDeploymentVersionData{tv1.DeploymentVersion().GetBuildId(): { Status: enumspb.WORKER_DEPLOYMENT_VERSION_STATUS_CURRENT, }}, []string{}, tqTypeAct) - data = s.getTaskQueueDeploymentData(tv, tqTypeAct) + data = s.getTaskQueueDeploymentData(env, tv, tqTypeAct) s.ProtoEqual(&persistencespb.DeploymentData{DeploymentsData: map[string]*persistencespb.WorkerDeploymentData{ tv1.DeploymentVersion().GetDeploymentName(): { RoutingConfig: routingConfig, @@ -3409,14 +3352,14 @@ func (s *Versioning3Suite) TestSyncDeploymentUserDataWithRoutingConfig_Update() CurrentVersionChangedTime: timestamp.TimePtr(t2), RevisionNumber: revisionNumber, } - s.syncTaskQueueDeploymentDataWithRoutingConfig(tv2, routingConfig, map[string]*deploymentspb.WorkerDeploymentVersionData{tv2.DeploymentVersion().GetBuildId(): { + s.syncTaskQueueDeploymentDataWithRoutingConfig(env, tv2, routingConfig, map[string]*deploymentspb.WorkerDeploymentVersionData{tv2.DeploymentVersion().GetBuildId(): { Status: enumspb.WORKER_DEPLOYMENT_VERSION_STATUS_CURRENT, }, tv1.DeploymentVersion().GetBuildId(): { Status: enumspb.WORKER_DEPLOYMENT_VERSION_STATUS_DRAINING, }}, []string{}, tqTypeAct, tqTypeWf) // Verify the data is right for the activity TQ. - data = s.getTaskQueueDeploymentData(tv, tqTypeAct) + data = s.getTaskQueueDeploymentData(env, tv, tqTypeAct) s.ProtoEqual(&persistencespb.DeploymentData{DeploymentsData: map[string]*persistencespb.WorkerDeploymentData{ tv1.DeploymentVersion().GetDeploymentName(): { RoutingConfig: routingConfig, @@ -3429,7 +3372,7 @@ func (s *Versioning3Suite) TestSyncDeploymentUserDataWithRoutingConfig_Update() }}, data) // Verify the data is right for the workflow TQ. - data = s.getTaskQueueDeploymentData(tv, tqTypeWf) + data = s.getTaskQueueDeploymentData(env, tv, tqTypeWf) s.ProtoEqual(&persistencespb.DeploymentData{DeploymentsData: map[string]*persistencespb.WorkerDeploymentData{ tv2.DeploymentVersion().GetDeploymentName(): { RoutingConfig: routingConfig, @@ -3442,8 +3385,8 @@ func (s *Versioning3Suite) TestSyncDeploymentUserDataWithRoutingConfig_Update() }}, data) // Forget v1 from both tq types - s.forgetDeploymentVersionsFromDeploymentData(tv1, tv1.DeploymentVersion().GetDeploymentName(), false, revisionNumber, tqTypeAct, tqTypeWf) - data = s.getTaskQueueDeploymentData(tv, tqTypeAct) + s.forgetDeploymentVersionsFromDeploymentData(env, tv1, tv1.DeploymentVersion().GetDeploymentName(), false, revisionNumber, tqTypeAct, tqTypeWf) + data = s.getTaskQueueDeploymentData(env, tv, tqTypeAct) s.ProtoEqual(&persistencespb.DeploymentData{DeploymentsData: map[string]*persistencespb.WorkerDeploymentData{ tv2.DeploymentVersion().GetDeploymentName(): { RoutingConfig: routingConfig, @@ -3453,7 +3396,7 @@ func (s *Versioning3Suite) TestSyncDeploymentUserDataWithRoutingConfig_Update() }, }}, data) - data = s.getTaskQueueDeploymentData(tv, tqTypeWf) + data = s.getTaskQueueDeploymentData(env, tv, tqTypeWf) s.ProtoEqual(&persistencespb.DeploymentData{DeploymentsData: map[string]*persistencespb.WorkerDeploymentData{ tv2.DeploymentVersion().GetDeploymentName(): { RoutingConfig: routingConfig, @@ -3464,8 +3407,8 @@ func (s *Versioning3Suite) TestSyncDeploymentUserDataWithRoutingConfig_Update() }}, data) // Forget v1 again should be a noop - s.forgetDeploymentVersionsFromDeploymentData(tv1, tv1.DeploymentVersion().GetDeploymentName(), false, revisionNumber, tqTypeAct, tqTypeWf) - data = s.getTaskQueueDeploymentData(tv, tqTypeAct) + s.forgetDeploymentVersionsFromDeploymentData(env, tv1, tv1.DeploymentVersion().GetDeploymentName(), false, revisionNumber, tqTypeAct, tqTypeWf) + data = s.getTaskQueueDeploymentData(env, tv, tqTypeAct) s.ProtoEqual(&persistencespb.DeploymentData{DeploymentsData: map[string]*persistencespb.WorkerDeploymentData{ tv2.DeploymentVersion().GetDeploymentName(): { RoutingConfig: routingConfig, @@ -3487,12 +3430,12 @@ func (s *Versioning3Suite) TestSyncDeploymentUserDataWithRoutingConfig_Update() RampingVersionPercentageChangedTime: timestamp.TimePtr(t3), RevisionNumber: revisionNumber, } - s.syncTaskQueueDeploymentDataWithRoutingConfig(tv2, routingConfig, map[string]*deploymentspb.WorkerDeploymentVersionData{ + s.syncTaskQueueDeploymentDataWithRoutingConfig(env, tv2, routingConfig, map[string]*deploymentspb.WorkerDeploymentVersionData{ tv2.DeploymentVersion().GetBuildId(): { Status: enumspb.WORKER_DEPLOYMENT_VERSION_STATUS_CURRENT, }, }, []string{}, tqTypeAct, tqTypeWf) - data = s.getTaskQueueDeploymentData(tv, tqTypeAct) + data = s.getTaskQueueDeploymentData(env, tv, tqTypeAct) s.ProtoEqual(&persistencespb.DeploymentData{DeploymentsData: map[string]*persistencespb.WorkerDeploymentData{ tv1.DeploymentVersion().GetDeploymentName(): { RoutingConfig: routingConfig, @@ -3504,47 +3447,45 @@ func (s *Versioning3Suite) TestSyncDeploymentUserDataWithRoutingConfig_Update() } -func (s *Versioning3Suite) setCurrentDeployment(tv *testvars.TestVars) { - ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second) - defer cancel() - s.Eventually(func() bool { +func (s *Versioning3Suite) setCurrentDeployment(env *testcore.TestEnv, tv *testvars.TestVars) { + failedPrecondition := serviceerror.NewFailedPreconditionf(workerdeployment.ErrCurrentVersionDoesNotHaveAllTaskQueues, tv.DeploymentVersionStringV32()).Error() + s.EventuallyWithT(func(t *assert.CollectT) { req := &workflowservice.SetWorkerDeploymentCurrentVersionRequest{ - Namespace: s.Namespace().String(), + Namespace: env.Namespace().String(), DeploymentName: tv.DeploymentSeries(), } req.BuildId = tv.BuildID() - _, err := s.FrontendClient().SetWorkerDeploymentCurrentVersion(ctx, req) - var notFound *serviceerror.NotFound - if errors.As(err, ¬Found) || (err != nil && strings.Contains(err.Error(), serviceerror.NewFailedPreconditionf(workerdeployment.ErrCurrentVersionDoesNotHaveAllTaskQueues, tv.DeploymentVersionStringV32()).Error())) { - return false + _, err := env.FrontendClient().SetWorkerDeploymentCurrentVersion(env.Context(), req) + if s.shouldRetryWorkerDeploymentRPC(env, err, failedPrecondition) { + require.NoError(t, err) + return } - s.NoError(err) - return err == nil + require.NoError(t, err) }, 60*time.Second, 500*time.Millisecond) // Wait for propagation to complete since we have tests using async entity workflows to set the current version - s.waitForDeploymentDataPropagationQueryWorkerDeployment(tv) + s.waitForDeploymentDataPropagationQueryWorkerDeployment(env, tv) } // pollUntilRegistered registers versioned pollers for the given deployment. // tqTypes controls which task queue types to poll; it defaults to workflow only. // Pollers run continuously until all TQ types are registered. -func (s *Versioning3Suite) pollUntilRegistered(ctx context.Context, tv *testvars.TestVars, tqTypes ...enumspb.TaskQueueType) { +func (s *Versioning3Suite) pollUntilRegistered(env *testcore.TestEnv, tv *testvars.TestVars, tqTypes ...enumspb.TaskQueueType) { if len(tqTypes) == 0 { tqTypes = []enumspb.TaskQueueType{tqTypeWf} } - pollCtx, cancel := context.WithCancel(ctx) + pollCtx, cancel := context.WithCancel(env.Context()) for _, tqType := range tqTypes { tqType := tqType go func() { for pollCtx.Err() == nil { switch tqType { case tqTypeWf: - s.idlePollWorkflow(pollCtx, tv, true, ver3MinPollTime, "should not get any tasks yet") + s.idlePollWorkflow(env, pollCtx, tv, true, ver3MinPollTime, "should not get any tasks yet") case tqTypeAct: - s.idlePollActivity(pollCtx, tv, true, ver3MinPollTime, "should not get any tasks yet") + s.idlePollActivity(env, tv, true, ver3MinPollTime, "should not get any tasks yet") case tqTypeNexus: - s.idlePollNexus(pollCtx, tv, true, ver3MinPollTime, "should not get any tasks yet") + s.idlePollNexus(env, pollCtx, tv, true, ver3MinPollTime, "should not get any tasks yet") default: panic("invalid task queue type") } @@ -3552,16 +3493,17 @@ func (s *Versioning3Suite) pollUntilRegistered(ctx context.Context, tv *testvars }() } // Wait until the version is visible and all requested task queue types are registered. - s.Eventually(func() bool { - resp, err := s.FrontendClient().DescribeWorkerDeploymentVersion(ctx, &workflowservice.DescribeWorkerDeploymentVersionRequest{ - Namespace: s.Namespace().String(), + s.EventuallyWithT(func(t *assert.CollectT) { + resp, err := env.FrontendClient().DescribeWorkerDeploymentVersion(env.Context(), &workflowservice.DescribeWorkerDeploymentVersionRequest{ + Namespace: env.Namespace().String(), Version: tv.DeploymentVersionString(), }) var notFound *serviceerror.NotFound if errors.As(err, ¬Found) { - return false + require.NoError(t, err) + return } - s.NoError(err) + require.NoError(t, err) tqName := tv.TaskQueue().GetName() for _, tqType := range tqTypes { found := false @@ -3571,82 +3513,98 @@ func (s *Versioning3Suite) pollUntilRegistered(ctx context.Context, tv *testvars break } } - if !found { - return false - } + require.True(t, found) } - return true }, 30*time.Second, 100*time.Millisecond) cancel() } -func (s *Versioning3Suite) unsetCurrentDeployment(tv *testvars.TestVars) { - ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second) - defer cancel() - s.Eventually(func() bool { +func (s *Versioning3Suite) unsetCurrentDeployment(env *testcore.TestEnv, tv *testvars.TestVars) { + s.EventuallyWithT(func(t *assert.CollectT) { req := &workflowservice.SetWorkerDeploymentCurrentVersionRequest{ - Namespace: s.Namespace().String(), + Namespace: env.Namespace().String(), DeploymentName: tv.DeploymentSeries(), } - _, err := s.FrontendClient().SetWorkerDeploymentCurrentVersion(ctx, req) - var notFound *serviceerror.NotFound - if errors.As(err, ¬Found) { - return false + _, err := env.FrontendClient().SetWorkerDeploymentCurrentVersion(env.Context(), req) + if s.shouldRetryWorkerDeploymentRPC(env, err) { + require.NoError(t, err) + return } - s.NoError(err) - return err == nil + require.NoError(t, err) }, 60*time.Second, 500*time.Millisecond) // Wait for propagation to complete since we have tests using async entity workflows to set the current version - s.waitForDeploymentDataPropagationQueryWorkerDeployment(tv) + s.waitForDeploymentDataPropagationQueryWorkerDeployment(env, tv) + s.waitForDeploymentDataPropagation(env, tv, versionStatusDraining, false, tqTypeWf) } func (s *Versioning3Suite) setRampingDeployment( + env *testcore.TestEnv, tv *testvars.TestVars, percentage float32, rampUnversioned bool, ) { - ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second) - defer cancel() bid := tv.BuildID() if rampUnversioned { bid = "" } + failedPrecondition := serviceerror.NewFailedPreconditionf(workerdeployment.ErrRampingVersionDoesNotHaveAllTaskQueues, tv.DeploymentVersionStringV32()).Error() - s.Eventually(func() bool { + s.EventuallyWithT(func(t *assert.CollectT) { req := &workflowservice.SetWorkerDeploymentRampingVersionRequest{ - Namespace: s.Namespace().String(), + Namespace: env.Namespace().String(), DeploymentName: tv.DeploymentSeries(), Percentage: percentage, } req.BuildId = bid - _, err := s.FrontendClient().SetWorkerDeploymentRampingVersion(ctx, req) - var notFound *serviceerror.NotFound - if errors.As(err, ¬Found) || (err != nil && strings.Contains(err.Error(), serviceerror.NewFailedPreconditionf(workerdeployment.ErrRampingVersionDoesNotHaveAllTaskQueues, tv.DeploymentVersionStringV32()).Error())) { - return false + _, err := env.FrontendClient().SetWorkerDeploymentRampingVersion(env.Context(), req) + if s.shouldRetryWorkerDeploymentRPC(env, err, failedPrecondition) { + require.NoError(t, err) + return } - s.NoError(err) - return err == nil + require.NoError(t, err) }, 60*time.Second, 500*time.Millisecond) // Wait for propagation to complete since we have tests using async entity workflows to set the current version - s.waitForDeploymentDataPropagationQueryWorkerDeployment(tv) + s.waitForDeploymentDataPropagationQueryWorkerDeployment(env, tv) } -func (s *Versioning3Suite) waitForDeploymentDataPropagationQueryWorkerDeployment(tv *testvars.TestVars) { - if s.deploymentWorkflowVersion == workerdeployment.AsyncSetCurrentAndRamping { - s.Eventually(func() bool { - resp, err := s.FrontendClient().DescribeWorkerDeployment(context.Background(), &workflowservice.DescribeWorkerDeploymentRequest{ - Namespace: s.Namespace().String(), +func (s *Versioning3Suite) waitForDeploymentDataPropagationQueryWorkerDeployment(env *testcore.TestEnv, tv *testvars.TestVars) { + if versioning3DeploymentWorkflowVersion == workerdeployment.AsyncSetCurrentAndRamping { + s.EventuallyWithT(func(t *assert.CollectT) { + resp, err := env.FrontendClient().DescribeWorkerDeployment(env.Context(), &workflowservice.DescribeWorkerDeploymentRequest{ + Namespace: env.Namespace().String(), DeploymentName: tv.DeploymentSeries(), }) - s.NoError(err) - return resp.GetWorkerDeploymentInfo().GetRoutingConfigUpdateState() == enumspb.ROUTING_CONFIG_UPDATE_STATE_COMPLETED + if s.shouldRetryWorkerDeploymentRPC(env, err) { + require.NoError(t, err) + return + } + require.NoError(t, err) + require.Equal(t, enumspb.ROUTING_CONFIG_UPDATE_STATE_COMPLETED, resp.GetWorkerDeploymentInfo().GetRoutingConfigUpdateState()) }, 10*time.Second, 100*time.Millisecond) } } +func (s *Versioning3Suite) shouldRetryWorkerDeploymentRPC(env *testcore.TestEnv, err error, retryableMessages ...string) bool { + if err == nil || env.Context().Err() != nil { + return false + } + var notFound *serviceerror.NotFound + if errors.As(err, ¬Found) || errors.Is(err, context.DeadlineExceeded) || common.IsRetryableRPCError(err) { + return true + } + errMsg := err.Error() + for _, msg := range retryableMessages { + if strings.Contains(errMsg, msg) { + return true + } + } + return false +} + func (s *Versioning3Suite) updateTaskQueueDeploymentData( + env *testcore.TestEnv, tv *testvars.TestVars, isCurrent bool, ramp float32, @@ -3654,7 +3612,7 @@ func (s *Versioning3Suite) updateTaskQueueDeploymentData( timeSinceUpdate time.Duration, tqTypes ...enumspb.TaskQueueType, ) { - s.syncTaskQueueDeploymentData(tv, isCurrent, ramp, rampUnversioned, time.Now().Add(-timeSinceUpdate), tqTypes...) + s.syncTaskQueueDeploymentData(env, tv, isCurrent, ramp, rampUnversioned, time.Now().Add(-timeSinceUpdate), tqTypes...) var status versionStatus if isCurrent { status = versionStatusCurrent @@ -3667,7 +3625,7 @@ func (s *Versioning3Suite) updateTaskQueueDeploymentData( status = versionStatusNil } - s.waitForDeploymentDataPropagation(tv, status, rampUnversioned, tqTypes...) + s.waitForDeploymentDataPropagation(env, tv, status, rampUnversioned, tqTypes...) } @@ -3675,6 +3633,7 @@ func (s *Versioning3Suite) updateTaskQueueDeploymentData( // and also waits for the data to propagate to all the relevant partitions. // TODO (Shivam): Update the name of this one. func (s *Versioning3Suite) updateTaskQueueDeploymentDataWithRoutingConfig( + env *testcore.TestEnv, tv *testvars.TestVars, routingConfig *deploymentpb.RoutingConfig, upsertVersions map[string]*deploymentspb.WorkerDeploymentVersionData, @@ -3682,16 +3641,16 @@ func (s *Versioning3Suite) updateTaskQueueDeploymentDataWithRoutingConfig( tqTypes ...enumspb.TaskQueueType, ) { - s.syncTaskQueueDeploymentDataWithRoutingConfig(tv, routingConfig, upsertVersions, forgetVersions, tqTypes...) + s.syncTaskQueueDeploymentDataWithRoutingConfig(env, tv, routingConfig, upsertVersions, forgetVersions, tqTypes...) // We need to know what the status of the version we are adding/forgetting is so that we can wait for it to propagate. for _, version := range upsertVersions { if version.GetStatus() == enumspb.WORKER_DEPLOYMENT_VERSION_STATUS_CURRENT { - s.waitForDeploymentDataPropagation(tv, versionStatusCurrent, false, tqTypes...) + s.waitForDeploymentDataPropagation(env, tv, versionStatusCurrent, false, tqTypes...) } else if version.GetStatus() == enumspb.WORKER_DEPLOYMENT_VERSION_STATUS_RAMPING { - s.waitForDeploymentDataPropagation(tv, versionStatusRamping, false, tqTypes...) + s.waitForDeploymentDataPropagation(env, tv, versionStatusRamping, false, tqTypes...) } else if version.GetStatus() == enumspb.WORKER_DEPLOYMENT_VERSION_STATUS_INACTIVE { - s.waitForDeploymentDataPropagation(tv, versionStatusInactive, false, tqTypes...) + s.waitForDeploymentDataPropagation(env, tv, versionStatusInactive, false, tqTypes...) } } } @@ -3700,15 +3659,14 @@ func (s *Versioning3Suite) updateTaskQueueDeploymentDataWithRoutingConfig( // returned from the WF type root partition, so no need to wait for propagation before calling this // function. func (s *Versioning3Suite) getTaskQueueDeploymentData( + env *testcore.TestEnv, tv *testvars.TestVars, tqType enumspb.TaskQueueType, ) *persistencespb.DeploymentData { - ctx, cancel := context.WithTimeout(context.Background(), time.Second*5) - defer cancel() - resp, err := s.GetTestCluster().MatchingClient().GetTaskQueueUserData( - ctx, + resp, err := env.GetTestCluster().MatchingClient().GetTaskQueueUserData( + env.Context(), &matchingservice.GetTaskQueueUserDataRequest{ - NamespaceId: s.NamespaceID().String(), + NamespaceId: env.NamespaceID().String(), TaskQueue: tv.TaskQueue().GetName(), TaskQueueType: tqTypeWf, }) @@ -3717,29 +3675,26 @@ func (s *Versioning3Suite) getTaskQueueDeploymentData( } func (s *Versioning3Suite) syncTaskQueueDeploymentDataWithRoutingConfig( + env *testcore.TestEnv, tv *testvars.TestVars, routingConfig *deploymentpb.RoutingConfig, upsertVersions map[string]*deploymentspb.WorkerDeploymentVersionData, forgetVersions []string, t ...enumspb.TaskQueueType, ) { - ctx, cancel := context.WithTimeout(context.Background(), time.Second*5) - defer cancel() - deploymentName := tv.DeploymentVersion().GetDeploymentName() var err error - _, err = s.GetTestCluster().MatchingClient().SyncDeploymentUserData( - ctx, &matchingservice.SyncDeploymentUserDataRequest{ - NamespaceId: s.NamespaceID().String(), + _, err = env.GetTestCluster().MatchingClient().SyncDeploymentUserData( + env.Context(), &matchingservice.SyncDeploymentUserDataRequest{ + NamespaceId: env.NamespaceID().String(), TaskQueue: tv.TaskQueue().GetName(), TaskQueueTypes: t, DeploymentName: deploymentName, UpdateRoutingConfig: routingConfig, UpsertVersionsData: upsertVersions, ForgetVersions: forgetVersions, - }, - ) + }) s.NoError(err) } @@ -3748,10 +3703,11 @@ func (s *Versioning3Suite) syncTaskQueueDeploymentDataWithRoutingConfig( // to an older version with revision number 0. This is used to test that workflows correctly use // inherited revision numbers instead of falling back to the (stale) current task queue version. func (s *Versioning3Suite) rollbackTaskQueueToVersion( + env *testcore.TestEnv, tv *testvars.TestVars, ) { - cleanup := s.InjectHook(testhooks.NewHook(testhooks.MatchingIgnoreRoutingConfigRevisionCheck, true)) + cleanup := env.InjectHook(testhooks.NewHook(testhooks.MatchingIgnoreRoutingConfigRevisionCheck, true)) defer cleanup() rc := &deploymentpb.RoutingConfig{ @@ -3759,24 +3715,26 @@ func (s *Versioning3Suite) rollbackTaskQueueToVersion( CurrentVersionChangedTime: timestamp.TimePtr(time.Now().Add(1 * time.Minute)), RevisionNumber: 0, } - s.syncTaskQueueDeploymentDataWithRoutingConfig(tv, rc, map[string]*deploymentspb.WorkerDeploymentVersionData{tv.DeploymentVersion().GetBuildId(): { + s.syncTaskQueueDeploymentDataWithRoutingConfig(env, tv, rc, map[string]*deploymentspb.WorkerDeploymentVersionData{tv.DeploymentVersion().GetBuildId(): { Status: enumspb.WORKER_DEPLOYMENT_VERSION_STATUS_CURRENT, }}, nil, tqTypeWf) // Verify that the rollback propagated to all partitions - s.Eventually(func() bool { - ms, err := s.GetTestCluster().MatchingClient().GetTaskQueueUserData(context.Background(), &matchingservice.GetTaskQueueUserDataRequest{ - NamespaceId: s.NamespaceID().String(), + s.EventuallyWithT(func(t *assert.CollectT) { + ms, err := env.GetTestCluster().MatchingClient().GetTaskQueueUserData(env.Context(), &matchingservice.GetTaskQueueUserDataRequest{ + NamespaceId: env.NamespaceID().String(), TaskQueue: tv.TaskQueue().GetName(), TaskQueueType: tqTypeWf, }) - s.NoError(err) + require.NoError(t, err) current, currentRevisionNumber, _, _, _, _, _, _ := worker_versioning.CalculateTaskQueueVersioningInfo(ms.GetUserData().GetData().GetPerType()[int32(tqTypeWf)].GetDeploymentData()) - return current.GetBuildId() == tv.DeploymentVersion().GetBuildId() && currentRevisionNumber == 0 + require.Equal(t, tv.DeploymentVersion().GetBuildId(), current.GetBuildId()) + require.Equal(t, int64(0), currentRevisionNumber) }, 10*time.Second, 100*time.Millisecond) } func (s *Versioning3Suite) syncTaskQueueDeploymentData( + env *testcore.TestEnv, tv *testvars.TestVars, isCurrent bool, ramp float32, @@ -3784,8 +3742,6 @@ func (s *Versioning3Suite) syncTaskQueueDeploymentData( updateTime time.Time, t ...enumspb.TaskQueueType, ) { - ctx, cancel := context.WithTimeout(context.Background(), time.Second*5) - defer cancel() v := tv.DeploymentVersion() if rampUnversioned { v = nil @@ -3800,9 +3756,9 @@ func (s *Versioning3Suite) syncTaskQueueDeploymentData( rampingSinceTime = routingUpdateTime } - _, err := s.GetTestCluster().MatchingClient().SyncDeploymentUserData( - ctx, &matchingservice.SyncDeploymentUserDataRequest{ - NamespaceId: s.NamespaceID().String(), + _, err := env.GetTestCluster().MatchingClient().SyncDeploymentUserData( + env.Context(), &matchingservice.SyncDeploymentUserDataRequest{ + NamespaceId: env.NamespaceID().String(), TaskQueue: tv.TaskQueue().GetName(), TaskQueueTypes: t, Operation: &matchingservice.SyncDeploymentUserDataRequest_UpdateVersionData{ @@ -3814,61 +3770,56 @@ func (s *Versioning3Suite) syncTaskQueueDeploymentData( RampPercentage: ramp, }, }, - }, - ) + }) s.NoError(err) } func (s *Versioning3Suite) forgetDeploymentVersionsFromDeploymentData( + env *testcore.TestEnv, tv *testvars.TestVars, deploymentName string, forgetUnversionedRamp bool, revisionNumber int64, t ...enumspb.TaskQueueType, ) { - ctx, cancel := context.WithTimeout(context.Background(), time.Second*5) - defer cancel() v := tv.DeploymentVersion() if forgetUnversionedRamp { v.BuildId = "" } - _, err := s.GetTestCluster().MatchingClient().SyncDeploymentUserData( - ctx, &matchingservice.SyncDeploymentUserDataRequest{ - NamespaceId: s.NamespaceID().String(), + _, err := env.GetTestCluster().MatchingClient().SyncDeploymentUserData( + env.Context(), &matchingservice.SyncDeploymentUserDataRequest{ + NamespaceId: env.NamespaceID().String(), TaskQueue: tv.TaskQueue().GetName(), TaskQueueTypes: t, DeploymentName: deploymentName, ForgetVersions: []string{tv.BuildID()}, - }, - ) + }) s.NoError(err) } func (s *Versioning3Suite) forgetTaskQueueDeploymentVersion( + env *testcore.TestEnv, tv *testvars.TestVars, t enumspb.TaskQueueType, forgetUnversionedRamp bool, ) { - ctx, cancel := context.WithTimeout(context.Background(), time.Second*5) - defer cancel() v := tv.DeploymentVersion() if forgetUnversionedRamp { v.BuildId = "" } - _, err := s.GetTestCluster().MatchingClient().SyncDeploymentUserData( - ctx, &matchingservice.SyncDeploymentUserDataRequest{ - NamespaceId: s.NamespaceID().String(), + _, err := env.GetTestCluster().MatchingClient().SyncDeploymentUserData( + env.Context(), &matchingservice.SyncDeploymentUserDataRequest{ + NamespaceId: env.NamespaceID().String(), TaskQueue: tv.TaskQueue().GetName(), TaskQueueTypes: []enumspb.TaskQueueType{t}, Operation: &matchingservice.SyncDeploymentUserDataRequest_ForgetVersion{ ForgetVersion: v, }, - }, - ) + }) s.NoError(err) } -func (s *Versioning3Suite) verifyWorkflowVersioning( +func (s *Versioning3Suite) verifyWorkflowVersioning(env *testcore.TestEnv, a *require.Assertions, tv *testvars.TestVars, behavior enumspb.VersioningBehavior, @@ -3876,14 +3827,13 @@ func (s *Versioning3Suite) verifyWorkflowVersioning( override *workflowpb.VersioningOverride, transition *workflowpb.DeploymentVersionTransition, ) { - dwf, err := s.FrontendClient().DescribeWorkflowExecution( - context.Background(), &workflowservice.DescribeWorkflowExecutionRequest{ - Namespace: s.Namespace().String(), + dwf, err := env.FrontendClient().DescribeWorkflowExecution( + env.Context(), &workflowservice.DescribeWorkflowExecutionRequest{ + Namespace: env.Namespace().String(), Execution: &commonpb.WorkflowExecution{ WorkflowId: tv.WorkflowID(), }, - }, - ) + }) a.NoError(err) versioningInfo := dwf.WorkflowExecutionInfo.GetVersioningInfo() @@ -4035,12 +3985,13 @@ func respondCompleteWorkflowUnversioned( } func (s *Versioning3Suite) startWorkflow( + env *testcore.TestEnv, tv *testvars.TestVars, override *workflowpb.VersioningOverride, ) string { request := &workflowservice.StartWorkflowExecutionRequest{ RequestId: tv.Any().String(), - Namespace: s.Namespace().String(), + Namespace: env.Namespace().String(), WorkflowId: tv.WorkflowID(), WorkflowType: tv.WorkflowType(), TaskQueue: tv.TaskQueue(), @@ -4048,72 +3999,62 @@ func (s *Versioning3Suite) startWorkflow( VersioningOverride: override, } - we, err0 := s.FrontendClient().StartWorkflowExecution(testcore.NewContext(), request) + we, err0 := env.FrontendClient().StartWorkflowExecution(testcore.NewContext(), request) s.NoError(err0) return we.GetRunId() } func (s *Versioning3Suite) queryWorkflow( + env *testcore.TestEnv, tv *testvars.TestVars, ) (*workflowservice.QueryWorkflowResponse, error) { request := &workflowservice.QueryWorkflowRequest{ - Namespace: s.Namespace().String(), + Namespace: env.Namespace().String(), Execution: tv.WorkflowExecution(), Query: tv.Query(), } shortCtx, cancel := context.WithTimeout(testcore.NewContext(), common.MinLongPollTimeout) defer cancel() - response, err := s.FrontendClient().QueryWorkflow(shortCtx, request) + response, err := env.FrontendClient().QueryWorkflow(shortCtx, request) return response, err } -// Name is used by testvars. We use a shorten test name in variables so that physical task queue IDs -// do not grow larger that DB column limit (currently as low as 272 chars). -func (s *Versioning3Suite) Name() string { - fullName := s.T().Name() - if len(fullName) <= 30 { - return fullName - } - short := fmt.Sprintf("%s-%08x", - fullName[len(fullName)-21:], - farm.Fingerprint32([]byte(fullName)), - ) - return strings.Replace(short, ".", "|", -1) -} - // pollWftAndHandle can be used in sync and async mode. For async mode pass the async channel. It // will be closed when the task is handled. // Returns the poller and poll response only in sync mode (can be used to process new wft in the response) func (s *Versioning3Suite) pollWftAndHandle( + env *testcore.TestEnv, tv *testvars.TestVars, sticky bool, async chan<- struct{}, handler func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error), ) (*taskpoller.TaskPoller, *workflowservice.RespondWorkflowTaskCompletedResponse) { - return s.doPollWftAndHandle(tv, true, sticky, async, handler) + return s.doPollWftAndHandle(env, tv, true, sticky, async, handler) } func (s *Versioning3Suite) unversionedPollWftAndHandle( + env *testcore.TestEnv, tv *testvars.TestVars, sticky bool, async chan<- struct{}, handler func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error), ) (*taskpoller.TaskPoller, *workflowservice.RespondWorkflowTaskCompletedResponse) { - return s.doPollWftAndHandle(tv, false, sticky, async, handler) + return s.doPollWftAndHandle(env, tv, false, sticky, async, handler) } // doPollWftAndHandle can be used in sync and async mode. For async mode pass the async channel. It // will be closed when the task is handled. // Returns the poller and poll response only in sync mode (can be used to process new wft in the response) func (s *Versioning3Suite) doPollWftAndHandle( + env *testcore.TestEnv, tv *testvars.TestVars, versioned bool, sticky bool, async chan<- struct{}, handler func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error), ) (*taskpoller.TaskPoller, *workflowservice.RespondWorkflowTaskCompletedResponse) { - poller := taskpoller.New(s.T(), s.FrontendClient(), s.Namespace().String()) + poller := taskpoller.New(s.T(), env.FrontendClient(), env.Namespace().String()) f := func() (*workflowservice.RespondWorkflowTaskCompletedResponse, error) { tq := tv.TaskQueue() if sticky { @@ -4124,7 +4065,7 @@ func (s *Versioning3Suite) doPollWftAndHandle( DeploymentOptions: tv.WorkerDeploymentOptions(versioned), TaskQueue: tq, }, - ).HandleTask(tv, handler, taskpoller.WithTimeout(30*time.Second)) + ).HandleTask(tv, handler, taskpoller.WithTimeout(time.Minute)) } if async == nil { resp, err := f() @@ -4140,12 +4081,13 @@ func (s *Versioning3Suite) doPollWftAndHandle( } func (s *Versioning3Suite) pollWftAndHandleQueries( + env *testcore.TestEnv, tv *testvars.TestVars, sticky bool, async chan<- any, handler func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondQueryTaskCompletedRequest, error), ) (*taskpoller.TaskPoller, *workflowservice.RespondQueryTaskCompletedResponse) { - poller := taskpoller.New(s.T(), s.FrontendClient(), s.Namespace().String()) + poller := taskpoller.New(s.T(), env.FrontendClient(), env.Namespace().String()) f := func() (*workflowservice.RespondQueryTaskCompletedResponse, error) { tq := tv.TaskQueue() if sticky { @@ -4171,12 +4113,13 @@ func (s *Versioning3Suite) pollWftAndHandleQueries( } func (s *Versioning3Suite) pollNexusTaskAndHandle( + env *testcore.TestEnv, tv *testvars.TestVars, sticky bool, async chan<- any, handler func(task *workflowservice.PollNexusTaskQueueResponse) (*workflowservice.RespondNexusTaskCompletedRequest, error), ) (*taskpoller.TaskPoller, *workflowservice.RespondNexusTaskCompletedResponse) { - poller := taskpoller.New(s.T(), s.FrontendClient(), s.Namespace().String()) + poller := taskpoller.New(s.T(), env.FrontendClient(), env.Namespace().String()) f := func() (*workflowservice.RespondNexusTaskCompletedResponse, error) { tq := tv.TaskQueue() if sticky { @@ -4202,35 +4145,40 @@ func (s *Versioning3Suite) pollNexusTaskAndHandle( } func (s *Versioning3Suite) unversionedPollActivityAndHandle( + env *testcore.TestEnv, tv *testvars.TestVars, async chan<- struct{}, handler func(task *workflowservice.PollActivityTaskQueueResponse) (*workflowservice.RespondActivityTaskCompletedRequest, error), ) { - s.doPollActivityAndHandle(tv, false, async, handler) + s.doPollActivityAndHandle(env, tv, false, async, handler) } func (s *Versioning3Suite) pollActivityAndHandle( + env *testcore.TestEnv, tv *testvars.TestVars, async chan<- struct{}, handler func(task *workflowservice.PollActivityTaskQueueResponse) (*workflowservice.RespondActivityTaskCompletedRequest, error), ) { - s.doPollActivityAndHandle(tv, true, async, handler) + s.doPollActivityAndHandle(env, tv, true, async, handler) +} + +func (s *Versioning3Suite) pollActivityAndHandleErr( + env *testcore.TestEnv, + tv *testvars.TestVars, + handler func(task *workflowservice.PollActivityTaskQueueResponse) (*workflowservice.RespondActivityTaskCompletedRequest, error), +) error { + return s.doPollActivityAndHandleErr(env, tv, true, handler) } func (s *Versioning3Suite) doPollActivityAndHandle( + env *testcore.TestEnv, tv *testvars.TestVars, versioned bool, async chan<- struct{}, handler func(task *workflowservice.PollActivityTaskQueueResponse) (*workflowservice.RespondActivityTaskCompletedRequest, error), ) { - poller := taskpoller.New(s.T(), s.FrontendClient(), s.Namespace().String()) f := func() error { - _, err := poller.PollActivityTask( - &workflowservice.PollActivityTaskQueueRequest{ - DeploymentOptions: tv.WorkerDeploymentOptions(versioned), - }, - ).HandleTask(tv, handler, taskpoller.WithTimeout(time.Minute)) - return err + return s.doPollActivityAndHandleErr(env, tv, versioned, handler) } if async == nil { s.NoError(f()) @@ -4242,25 +4190,35 @@ func (s *Versioning3Suite) doPollActivityAndHandle( } } -func (s *Versioning3Suite) idlePollWorkflow( +func (s *Versioning3Suite) doPollActivityAndHandleErr( + env *testcore.TestEnv, + tv *testvars.TestVars, + versioned bool, + handler func(task *workflowservice.PollActivityTaskQueueResponse) (*workflowservice.RespondActivityTaskCompletedRequest, error), +) error { + poller := taskpoller.New(s.T(), env.FrontendClient(), env.Namespace().String()) + _, err := poller.PollActivityTask( + &workflowservice.PollActivityTaskQueueRequest{ + DeploymentOptions: tv.WorkerDeploymentOptions(versioned), + }).HandleTask(tv, handler, taskpoller.WithTimeout(time.Minute)) + return err +} + +func (s *Versioning3Suite) idlePollWorkflow(env *testcore.TestEnv, ctx context.Context, tv *testvars.TestVars, versioned bool, timeout time.Duration, unexpectedTaskMessage string, ) { - poller := taskpoller.New(s.T(), s.FrontendClient(), s.Namespace().String()) + poller := taskpoller.New(s.T(), env.FrontendClient(), env.Namespace().String()) _, _ = poller.PollWorkflowTask( &workflowservice.PollWorkflowTaskQueueRequest{ DeploymentOptions: tv.WorkerDeploymentOptions(versioned), - }, - ).HandleTask( + }).HandleTask( tv, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { - // Can't call s.Fail directly because linter will complain about unsafe assertion in another go routine. - // s.Assert().Fail is safe but linter automatically converts it to s.Fail if written directly! - a := s.Assert() - a.Fail(unexpectedTaskMessage) + s.T().Error(unexpectedTaskMessage) return nil, nil }, taskpoller.WithTimeout(timeout), @@ -4269,18 +4227,18 @@ func (s *Versioning3Suite) idlePollWorkflow( } func (s *Versioning3Suite) idlePollUnversionedActivity( + env *testcore.TestEnv, tv *testvars.TestVars, timeout time.Duration, unexpectedTaskMessage string, ) { - poller := taskpoller.New(s.T(), s.FrontendClient(), s.Namespace().String()) + poller := taskpoller.New(s.T(), env.FrontendClient(), env.Namespace().String()) _, _ = poller.PollActivityTask( - &workflowservice.PollActivityTaskQueueRequest{}, - ).HandleTask( + &workflowservice.PollActivityTaskQueueRequest{}).HandleTask( tv, func(task *workflowservice.PollActivityTaskQueueResponse) (*workflowservice.RespondActivityTaskCompletedRequest, error) { if task != nil { - s.Logger.Error(fmt.Sprintf("Unexpected activity task received, ID: %s", task.ActivityId)) + env.Logger.Error(fmt.Sprintf("Unexpected activity task received, ID: %s", task.ActivityId)) s.Fail(unexpectedTaskMessage) } return nil, nil @@ -4290,40 +4248,38 @@ func (s *Versioning3Suite) idlePollUnversionedActivity( } func (s *Versioning3Suite) idlePollActivity( - ctx context.Context, + env *testcore.TestEnv, tv *testvars.TestVars, versioned bool, timeout time.Duration, unexpectedTaskMessage string, ) { - poller := taskpoller.New(s.T(), s.FrontendClient(), s.Namespace().String()) + poller := taskpoller.New(s.T(), env.FrontendClient(), env.Namespace().String()) _, _ = poller.PollActivityTask( &workflowservice.PollActivityTaskQueueRequest{ DeploymentOptions: tv.WorkerDeploymentOptions(versioned), - }, - ).HandleTask( + }).HandleTask( tv, func(task *workflowservice.PollActivityTaskQueueResponse) (*workflowservice.RespondActivityTaskCompletedRequest, error) { if task != nil { - s.Logger.Error(fmt.Sprintf("Unexpected activity task received, ID: %s", task.ActivityId)) - a := s.Assert() - a.Fail(unexpectedTaskMessage) + env.Logger.Error(fmt.Sprintf("Unexpected activity task received, ID: %s", task.ActivityId)) + s.T().Error(unexpectedTaskMessage) } return nil, nil }, taskpoller.WithTimeout(timeout), - taskpoller.WithContext(ctx), + taskpoller.WithContext(env.Context()), ) } -func (s *Versioning3Suite) idlePollNexus( +func (s *Versioning3Suite) idlePollNexus(env *testcore.TestEnv, ctx context.Context, tv *testvars.TestVars, versioned bool, timeout time.Duration, unexpectedTaskMessage string, ) { - poller := taskpoller.New(s.T(), s.FrontendClient(), s.Namespace().String()) + poller := taskpoller.New(s.T(), env.FrontendClient(), env.Namespace().String()) _, _ = poller.PollNexusTask( &workflowservice.PollNexusTaskQueueRequest{ DeploymentOptions: tv.WorkerDeploymentOptions(versioned), @@ -4331,8 +4287,7 @@ func (s *Versioning3Suite) idlePollNexus( tv, func(task *workflowservice.PollNexusTaskQueueResponse) (*workflowservice.RespondNexusTaskCompletedRequest, error) { if task != nil { - a := s.Assert() - a.Fail(unexpectedTaskMessage) + s.T().Error(unexpectedTaskMessage) } return nil, nil }, @@ -4342,14 +4297,14 @@ func (s *Versioning3Suite) idlePollNexus( } func (s *Versioning3Suite) verifyWorkflowStickyQueue( + env *testcore.TestEnv, tv *testvars.TestVars, ) { - ms, err := s.GetTestCluster().HistoryClient().GetMutableState( - context.Background(), &historyservice.GetMutableStateRequest{ - NamespaceId: s.NamespaceID().String(), + ms, err := env.GetTestCluster().HistoryClient().GetMutableState( + env.Context(), &historyservice.GetMutableStateRequest{ + NamespaceId: env.NamespaceID().String(), Execution: tv.WorkflowExecution(), - }, - ) + }) s.NoError(err) s.Equal(tv.StickyTaskQueue().GetName(), ms.StickyTaskQueue.GetName()) } @@ -4357,14 +4312,14 @@ func (s *Versioning3Suite) verifyWorkflowStickyQueue( // Sticky queue needs to be created in server before tasks can schedule in it. Call to this method // create the sticky queue by polling it. func (s *Versioning3Suite) warmUpSticky( + env *testcore.TestEnv, tv *testvars.TestVars, ) { - poller := taskpoller.New(s.T(), s.FrontendClient(), s.Namespace().String()) + poller := taskpoller.New(s.T(), env.FrontendClient(), env.Namespace().String()) _, _ = poller.PollWorkflowTask( &workflowservice.PollWorkflowTaskQueueRequest{ TaskQueue: tv.StickyTaskQueue(), - }, - ).HandleTask( + }).HandleTask( tv, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.Fail("sticky task is not expected") @@ -4376,15 +4331,13 @@ func (s *Versioning3Suite) warmUpSticky( // TODO (Shivam): Clean up this function once sync entity workflows have been removed. func (s *Versioning3Suite) waitForDeploymentDataPropagation( + env *testcore.TestEnv, tv *testvars.TestVars, status versionStatus, unversionedRamp bool, tqTypes ...enumspb.TaskQueueType, ) { - ctx, cancel := context.WithTimeout(context.Background(), time.Second*30) - defer cancel() - - v := s.GetTestCluster().Host().DcClient().GetValue(dynamicconfig.MatchingNumTaskqueueReadPartitions.Key()) + v := env.GetTestCluster().Host().DcClient().GetValue(dynamicconfig.MatchingNumTaskqueueReadPartitions.Key()) s.NotEmpty(v, "versioning tests require setting explicit number of partitions") count, ok := v[0].Value.(int) s.True(ok, "partition count is not an int") @@ -4400,7 +4353,7 @@ func (s *Versioning3Suite) waitForDeploymentDataPropagation( remaining[partAndType{i, tqt}] = struct{}{} } } - f, err := tqid.NewTaskQueueFamily(s.NamespaceID().String(), tv.TaskQueue().GetName()) + f, err := tqid.NewTaskQueueFamily(env.NamespaceID().String(), tv.TaskQueue().GetName()) s.EventuallyWithT(func(t *assert.CollectT) { a := require.New(t) for pt := range remaining { @@ -4408,10 +4361,10 @@ func (s *Versioning3Suite) waitForDeploymentDataPropagation( partition := f.TaskQueue(pt.tp).NormalPartition(pt.part) // Use lower-level GetTaskQueueUserData instead of GetWorkerBuildIdCompatibility // here so that we can target activity queues. - res, err := s.GetTestCluster().MatchingClient().GetTaskQueueUserData( - ctx, + res, err := env.GetTestCluster().MatchingClient().GetTaskQueueUserData( + env.Context(), &matchingservice.GetTaskQueueUserDataRequest{ - NamespaceId: s.NamespaceID().String(), + NamespaceId: env.NamespaceID().String(), TaskQueue: partition.RpcName(), TaskQueueType: partition.TaskType(), }) @@ -4471,19 +4424,17 @@ func (s *Versioning3Suite) waitForDeploymentDataPropagation( } func (s *Versioning3Suite) validateBacklogCount( + env *testcore.TestEnv, tv *testvars.TestVars, tqType enumspb.TaskQueueType, expectedCount int64, ) { - ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) - defer cancel() - var resp *workflowservice.DescribeTaskQueueResponse var err error s.EventuallyWithT(func(t *assert.CollectT) { - resp, err = s.FrontendClient().DescribeTaskQueue(ctx, &workflowservice.DescribeTaskQueueRequest{ - Namespace: s.Namespace().String(), + resp, err = env.FrontendClient().DescribeTaskQueue(env.Context(), &workflowservice.DescribeTaskQueueRequest{ + Namespace: env.Namespace().String(), TaskQueue: tv.TaskQueue(), TaskQueueType: tqType, ReportStats: true, @@ -4498,14 +4449,12 @@ func (s *Versioning3Suite) validateBacklogCount( } func (s *Versioning3Suite) verifyVersioningSAs( + env *testcore.TestEnv, tv *testvars.TestVars, behavior enumspb.VersioningBehavior, executionStatus enumspb.WorkflowExecutionStatus, usedBuilds ...*testvars.TestVars, ) { - ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) - defer cancel() - s.EventuallyWithT(func(t *assert.CollectT) { var query string if behavior != vbUnspecified { @@ -4515,8 +4464,8 @@ func (s *Versioning3Suite) verifyVersioningSAs( query = fmt.Sprintf("WorkflowId = '%s' AND TemporalWorkerDeploymentVersion is null AND TemporalWorkflowVersioningBehavior is null AND ExecutionStatus = '%s'", tv.WorkflowID(), executionStatus) } - resp, err := s.FrontendClient().ListWorkflowExecutions(ctx, &workflowservice.ListWorkflowExecutionsRequest{ - Namespace: s.Namespace().String(), + resp, err := env.FrontendClient().ListWorkflowExecutions(env.Context(), &workflowservice.ListWorkflowExecutionsRequest{ + Namespace: env.Namespace().String(), Query: query, }) a := assert.New(t) @@ -4559,10 +4508,12 @@ func (s *Versioning3Suite) verifyVersioningSAs( } func (s *Versioning3Suite) TestAutoUpgradeWorkflows_NoBouncingBetweenVersions() { - s.OverrideDynamicConfig(dynamicconfig.MatchingNumTaskqueueReadPartitions, 1) - s.OverrideDynamicConfig(dynamicconfig.MatchingNumTaskqueueWritePartitions, 1) + env := s.setupEnv( + testcore.WithDynamicConfig(dynamicconfig.MatchingNumTaskqueueReadPartitions, 1), + testcore.WithDynamicConfig(dynamicconfig.MatchingNumTaskqueueWritePartitions, 1), + ) - tv0 := testvars.New(s).WithBuildIDNumber(0) + tv0 := env.Tv().WithBuildIDNumber(0) tv1 := tv0.WithBuildIDNumber(1) wf := func(ctx workflow.Context) (string, error) { @@ -4571,7 +4522,7 @@ func (s *Versioning3Suite) TestAutoUpgradeWorkflows_NoBouncingBetweenVersions() } // v1 workers - w1 := worker.New(s.SdkClient(), tv1.TaskQueue().GetName(), worker.Options{ + w1 := worker.New(env.SdkClient(), tv1.TaskQueue().GetName(), worker.Options{ DeploymentOptions: worker.DeploymentOptions{Version: tv1.SDKDeploymentVersion(), UseVersioning: true, DefaultVersioningBehavior: workflow.VersioningBehaviorAutoUpgrade}, }) w1.RegisterWorkflowWithOptions(wf, workflow.RegisterOptions{Name: "wf", VersioningBehavior: workflow.VersioningBehaviorAutoUpgrade}) @@ -4579,10 +4530,10 @@ func (s *Versioning3Suite) TestAutoUpgradeWorkflows_NoBouncingBetweenVersions() defer w1.Stop() // Set v1 to be the current version - s.setCurrentDeployment(tv1) + s.setCurrentDeployment(env, tv1) // Start a workflow on v1 - run, err := s.SdkClient().ExecuteWorkflow(context.Background(), sdkclient.StartWorkflowOptions{ + run, err := env.SdkClient().ExecuteWorkflow(env.Context(), sdkclient.StartWorkflowOptions{ ID: tv1.WorkflowID(), TaskQueue: tv1.TaskQueue().GetName(), }, "wf") @@ -4591,27 +4542,32 @@ func (s *Versioning3Suite) TestAutoUpgradeWorkflows_NoBouncingBetweenVersions() // Verify that the workflow is running on v1 s.EventuallyWithT(func(t *assert.CollectT) { a := require.New(t) - s.verifyWorkflowVersioning(a, tv1, vbUnpinned, tv1.Deployment(), nil, nil) + s.verifyWorkflowVersioning(env, a, tv1, vbUnpinned, tv1.Deployment(), nil, nil) }, 10*time.Second, 100*time.Millisecond) // Start v0 workers to ensure they never receive a task - idlePollerCtx, idlePollerCancel := context.WithTimeout(context.Background(), 10*time.Second) + idlePollerCtx, idlePollerCancel := context.WithTimeout(env.Context(), 10*time.Second) defer idlePollerCancel() - go s.idlePollWorkflow(idlePollerCtx, tv0, true, ver3MinPollTime, "workflows should not go to the old deployment") + go s.idlePollWorkflow(env, idlePollerCtx, tv0, true, ver3MinPollTime, "workflows should not go to the old deployment") // Rollback the userData to v0 to simulate routing config lag within a single partition. - s.rollbackTaskQueueToVersion(tv0) + s.rollbackTaskQueueToVersion(env, tv0) // Unblock the workflow to let it continue - s.NoError(s.SdkClient().SignalWorkflow(context.Background(), tv1.WorkflowID(), "", "afterRollback", nil)) + s.NoError(env.SdkClient().SignalWorkflow(env.Context(), tv1.WorkflowID(), "", "afterRollback", nil)) // Verify that the workflow completed successfully on v1 var result string - s.NoError(run.Get(context.Background(), &result)) + s.NoError(run.Get(env.Context(), &result)) s.Equal("v1", result) } func (s *Versioning3Suite) TestWorkflowTQLags_DependentActivityStartsTransition() { + env := s.setupEnv( + testcore.WithDynamicConfig(dynamicconfig.MatchingNumTaskqueueReadPartitions, 1), + testcore.WithDynamicConfig(dynamicconfig.MatchingNumTaskqueueWritePartitions, 1), + testcore.WithDynamicConfig(dynamicconfig.UseRevisionNumberForWorkerVersioning, true), + ) /* The aim of this test is to show the following does not occur when using revisionNumber mechanics: - If the workflow TQ lags behind the activity TQ, with respect to the current version of a deployment, the activity should not be @@ -4627,15 +4583,11 @@ func (s *Versioning3Suite) TestWorkflowTQLags_DependentActivityStartsTransition( - Let an activity poller complete the activity task on v1. - We should see a workflow transition happen to v1 even though the workflow TQ is lagging behind the activity TQ. */ - s.OverrideDynamicConfig(dynamicconfig.MatchingNumTaskqueueReadPartitions, 1) - s.OverrideDynamicConfig(dynamicconfig.MatchingNumTaskqueueWritePartitions, 1) - s.OverrideDynamicConfig(dynamicconfig.UseRevisionNumberForWorkerVersioning, true) - - tv0 := testvars.New(s).WithBuildIDNumber(0) + tv0 := env.Tv().WithBuildIDNumber(0) tv1 := tv0.WithBuildIDNumber(1) // Update the userData by setting the current version to v0 - s.updateTaskQueueDeploymentDataWithRoutingConfig(tv0, &deploymentpb.RoutingConfig{ + s.updateTaskQueueDeploymentDataWithRoutingConfig(env, tv0, &deploymentpb.RoutingConfig{ CurrentDeploymentVersion: worker_versioning.ExternalWorkerDeploymentVersionFromStringV31(tv0.DeploymentVersionString()), CurrentVersionChangedTime: timestamp.TimePtr(time.Now()), RevisionNumber: 1, @@ -4644,13 +4596,13 @@ func (s *Versioning3Suite) TestWorkflowTQLags_DependentActivityStartsTransition( }}, []string{}, tqTypeWf, tqTypeAct) // Wait until all task queue partitions know that v0 is current. - s.waitForDeploymentDataPropagation(tv0, versionStatusCurrent, false, tqTypeWf, tqTypeAct) + s.waitForDeploymentDataPropagation(env, tv0, versionStatusCurrent, false, tqTypeWf, tqTypeAct) // Start a workflow on v0. - s.startWorkflow(tv0, nil) + s.startWorkflow(env, tv0, nil) // Poll for the workflow task on v0. - s.pollWftAndHandle(tv0, false, nil, + s.pollWftAndHandle(env, tv0, false, nil, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) return respondWftWithActivities(tv0, tv0, false, vbUnpinned, "activity1"), nil @@ -4659,11 +4611,11 @@ func (s *Versioning3Suite) TestWorkflowTQLags_DependentActivityStartsTransition( // Verify that the workflow is running on v1. s.EventuallyWithT(func(t *assert.CollectT) { a := require.New(t) - s.verifyWorkflowVersioning(a, tv0, vbUnpinned, tv0.Deployment(), nil, nil) + s.verifyWorkflowVersioning(env, a, tv0, vbUnpinned, tv0.Deployment(), nil, nil) }, 10*time.Second, 100*time.Millisecond) // Update the userData for the activity TQ by setting the current version to v1. - s.syncTaskQueueDeploymentDataWithRoutingConfig(tv1, &deploymentpb.RoutingConfig{ + s.syncTaskQueueDeploymentDataWithRoutingConfig(env, tv1, &deploymentpb.RoutingConfig{ CurrentDeploymentVersion: worker_versioning.ExternalWorkerDeploymentVersionFromStringV31(tv1.DeploymentVersionString()), CurrentVersionChangedTime: timestamp.TimePtr(time.Now()), RevisionNumber: 2, @@ -4672,21 +4624,21 @@ func (s *Versioning3Suite) TestWorkflowTQLags_DependentActivityStartsTransition( }}, []string{}, tqTypeAct) // Wait until all task queue partitions know that v1 is current. - s.waitForDeploymentDataPropagation(tv1, versionStatusCurrent, false, tqTypeAct) + s.waitForDeploymentDataPropagation(env, tv1, versionStatusCurrent, false, tqTypeAct) // Poll and complete the workflow task which should have been scheduled by the activity task since it would have started a transition to v1. workflowTaskCh := make(chan struct{}, 1) activityTaskCh := make(chan struct{}, 1) // Poll and complete the activity task on v1. - s.pollActivityAndHandle(tv1, activityTaskCh, + s.pollActivityAndHandle(env, tv1, activityTaskCh, func(task *workflowservice.PollActivityTaskQueueResponse) (*workflowservice.RespondActivityTaskCompletedRequest, error) { s.NotNil(task) return respondActivity(), nil }) // Workflow task poller transitions the workflow to v1. - s.pollWftAndHandle(tv1, false, workflowTaskCh, + s.pollWftAndHandle(env, tv1, false, workflowTaskCh, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) return respondEmptyWft(tv1, false, vbUnpinned), nil @@ -4698,11 +4650,16 @@ func (s *Versioning3Suite) TestWorkflowTQLags_DependentActivityStartsTransition( // Verify that the workflow is running on v1. s.EventuallyWithT(func(t *assert.CollectT) { a := require.New(t) - s.verifyWorkflowVersioning(a, tv0, vbUnpinned, tv1.Deployment(), nil, nil) + s.verifyWorkflowVersioning(env, a, tv0, vbUnpinned, tv1.Deployment(), nil, nil) }, 10*time.Second, 100*time.Millisecond) } func (s *Versioning3Suite) TestActivityTQLags_DependentActivityCompletesOnTheNewVersion() { + env := s.setupEnv( + testcore.WithDynamicConfig(dynamicconfig.MatchingNumTaskqueueReadPartitions, 1), + testcore.WithDynamicConfig(dynamicconfig.MatchingNumTaskqueueWritePartitions, 1), + testcore.WithDynamicConfig(dynamicconfig.UseRevisionNumberForWorkerVersioning, true), + ) /* The aim of this test is to show the following does not occur when using revisionNumber mechanics: - If the activity TQ lags behind the workflow TQ, with respect to the current version of a deployment, the activity should not be @@ -4718,18 +4675,11 @@ func (s *Versioning3Suite) TestActivityTQLags_DependentActivityCompletesOnTheNew - Let a controlled poller complete the workflow task on v1 and schedule an activity task. - We should see the activity task be dispatched to a v1 poller even though it's lagging behind the workflow TQ. */ - ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) - defer cancel() - - s.OverrideDynamicConfig(dynamicconfig.MatchingNumTaskqueueReadPartitions, 1) - s.OverrideDynamicConfig(dynamicconfig.MatchingNumTaskqueueWritePartitions, 1) - s.OverrideDynamicConfig(dynamicconfig.UseRevisionNumberForWorkerVersioning, true) - - tv0 := testvars.New(s).WithBuildIDNumber(0) + tv0 := env.Tv().WithBuildIDNumber(0) tv1 := tv0.WithBuildIDNumber(1) // Update the userData for the workflow TQ by setting the current version to v0 - s.updateTaskQueueDeploymentDataWithRoutingConfig(tv0, &deploymentpb.RoutingConfig{ + s.updateTaskQueueDeploymentDataWithRoutingConfig(env, tv0, &deploymentpb.RoutingConfig{ CurrentDeploymentVersion: worker_versioning.ExternalWorkerDeploymentVersionFromStringV31(tv0.DeploymentVersionString()), CurrentVersionChangedTime: timestamp.TimePtr(time.Now()), RevisionNumber: 1, @@ -4737,13 +4687,13 @@ func (s *Versioning3Suite) TestActivityTQLags_DependentActivityCompletesOnTheNew Status: enumspb.WORKER_DEPLOYMENT_VERSION_STATUS_CURRENT, }}, []string{}, tqTypeWf, tqTypeAct) // Wait until all task queue partitions know that v0 is current. - s.waitForDeploymentDataPropagation(tv0, versionStatusCurrent, false, tqTypeWf, tqTypeAct) + s.waitForDeploymentDataPropagation(env, tv0, versionStatusCurrent, false, tqTypeWf, tqTypeAct) // Start a workflow on v0. - s.startWorkflow(tv0, nil) + s.startWorkflow(env, tv0, nil) // Let a controlled poller complete the workflow task on v0 and schedule another workflow task. - s.pollWftAndHandle(tv0, false, nil, + s.pollWftAndHandle(env, tv0, false, nil, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) resp := respondEmptyWft(tv0, false, vbUnpinned) @@ -4754,11 +4704,11 @@ func (s *Versioning3Suite) TestActivityTQLags_DependentActivityCompletesOnTheNew // Verify that the workflow is running on v0. s.EventuallyWithT(func(t *assert.CollectT) { a := require.New(t) - s.verifyWorkflowVersioning(a, tv0, vbUnpinned, tv0.Deployment(), nil, nil) + s.verifyWorkflowVersioning(env, a, tv0, vbUnpinned, tv0.Deployment(), nil, nil) }, 10*time.Second, 100*time.Millisecond) // Update the userData for the workflow TQ *only* by setting the current version to v1 - s.updateTaskQueueDeploymentDataWithRoutingConfig(tv1, &deploymentpb.RoutingConfig{ + s.updateTaskQueueDeploymentDataWithRoutingConfig(env, tv1, &deploymentpb.RoutingConfig{ CurrentDeploymentVersion: worker_versioning.ExternalWorkerDeploymentVersionFromStringV31(tv1.DeploymentVersionString()), CurrentVersionChangedTime: timestamp.TimePtr(time.Now()), RevisionNumber: 2, @@ -4766,15 +4716,15 @@ func (s *Versioning3Suite) TestActivityTQLags_DependentActivityCompletesOnTheNew Status: enumspb.WORKER_DEPLOYMENT_VERSION_STATUS_CURRENT, }}, []string{}, tqTypeWf) // Wait until all task queue partitions know that v1 is current. - s.waitForDeploymentDataPropagation(tv1, versionStatusCurrent, false, tqTypeWf) + s.waitForDeploymentDataPropagation(env, tv1, versionStatusCurrent, false, tqTypeWf) // Register the v1 worker in the activity TQ to prevent matching from thinking this as an independent unpinned activity - s.updateTaskQueueDeploymentDataWithRoutingConfig(tv1, nil, map[string]*deploymentspb.WorkerDeploymentVersionData{tv1.DeploymentVersion().GetBuildId(): { + s.updateTaskQueueDeploymentDataWithRoutingConfig(env, tv1, nil, map[string]*deploymentspb.WorkerDeploymentVersionData{tv1.DeploymentVersion().GetBuildId(): { Status: enumspb.WORKER_DEPLOYMENT_VERSION_STATUS_INACTIVE, }}, []string{}, tqTypeAct) // Let a controlled poller complete the workflow task on v1 and schedule an activity task - s.pollWftAndHandle(tv1, false, nil, + s.pollWftAndHandle(env, tv1, false, nil, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) return respondWftWithActivities(tv1, tv1, false, vbUnpinned, "activity1"), nil @@ -4782,22 +4732,22 @@ func (s *Versioning3Suite) TestActivityTQLags_DependentActivityCompletesOnTheNew // Start an idle activity poller on v0. This poller should not receive any activity tasks //nolint:testifylint - go s.idlePollActivity(ctx, tv0, true, ver3MinPollTime, "activity should not go to the old deployment") + go s.idlePollActivity(env, tv0, true, ver3MinPollTime, "activity should not go to the old deployment") // Start a poller on v1 activityTaskCh := make(chan struct{}, 1) - s.pollActivityAndHandle(tv1, activityTaskCh, + s.pollActivityAndHandle(env, tv1, activityTaskCh, func(task *workflowservice.PollActivityTaskQueueResponse) (*workflowservice.RespondActivityTaskCompletedRequest, error) { s.NotNil(task) return respondActivity(), nil }) - s.WaitForChannel(ctx, activityTaskCh) + env.WaitForChannel(env.Context(), activityTaskCh) // Verify that the workflow is still running on v1. s.EventuallyWithT(func(t *assert.CollectT) { a := require.New(t) - s.verifyWorkflowVersioning(a, tv1, vbUnpinned, tv1.Deployment(), nil, nil) + s.verifyWorkflowVersioning(env, a, tv1, vbUnpinned, tv1.Deployment(), nil, nil) }, 10*time.Second, 100*time.Millisecond) } @@ -4805,10 +4755,12 @@ func (s *Versioning3Suite) TestActivityTQLags_DependentActivityCompletesOnTheNew // the test is present to show that revision number mechanics work as expected even when the task-queue // partitions have a more updated view of the current version than the mutable state of a workflow. func (s *Versioning3Suite) TestChildStartsWithParentRevision_SameTQ_TQAhead() { - s.OverrideDynamicConfig(dynamicconfig.MatchingNumTaskqueueReadPartitions, 1) - s.OverrideDynamicConfig(dynamicconfig.MatchingNumTaskqueueWritePartitions, 1) + env := s.setupEnv( + testcore.WithDynamicConfig(dynamicconfig.MatchingNumTaskqueueReadPartitions, 1), + testcore.WithDynamicConfig(dynamicconfig.MatchingNumTaskqueueWritePartitions, 1), + ) - tvParent := testvars.New(s).WithBuildIDNumber(1).WithWorkflowIDNumber(1) + tvParent := env.Tv().WithBuildIDNumber(1).WithWorkflowIDNumber(1) tvChild := tvParent.WithWorkflowIDNumber(2) // Parent workflow waits for a signal before starting the child. @@ -4831,7 +4783,7 @@ func (s *Versioning3Suite) TestChildStartsWithParentRevision_SameTQ_TQAhead() { } // Workers on parent and child TQs. Note: Both the parent and child workflows run on the same TQ. - w := worker.New(s.SdkClient(), tvParent.TaskQueue().GetName(), worker.Options{ + w := worker.New(env.SdkClient(), tvParent.TaskQueue().GetName(), worker.Options{ DeploymentOptions: worker.DeploymentOptions{Version: tvParent.SDKDeploymentVersion(), UseVersioning: true, DefaultVersioningBehavior: workflow.VersioningBehaviorAutoUpgrade}, }) w.RegisterWorkflowWithOptions(parent, workflow.RegisterOptions{Name: "parent", VersioningBehavior: workflow.VersioningBehaviorAutoUpgrade}) @@ -4841,7 +4793,7 @@ func (s *Versioning3Suite) TestChildStartsWithParentRevision_SameTQ_TQAhead() { // Also start a v2 worker, which polls on the same TQ. tv2 := tvParent.WithBuildIDNumber(2) - w2 := worker.New(s.SdkClient(), tv2.TaskQueue().GetName(), worker.Options{ + w2 := worker.New(env.SdkClient(), tv2.TaskQueue().GetName(), worker.Options{ DeploymentOptions: worker.DeploymentOptions{Version: tv2.SDKDeploymentVersion(), UseVersioning: true, DefaultVersioningBehavior: workflow.VersioningBehaviorAutoUpgrade}, }) w2.RegisterWorkflowWithOptions(parent, workflow.RegisterOptions{Name: "parent", VersioningBehavior: workflow.VersioningBehaviorAutoUpgrade}) @@ -4850,53 +4802,43 @@ func (s *Versioning3Suite) TestChildStartsWithParentRevision_SameTQ_TQAhead() { defer w2.Stop() // Set v1 as the current version for the deployment - s.setCurrentDeployment(tvParent) + s.setCurrentDeployment(env, tvParent) // Start parent - ctx, cancel := context.WithTimeout(context.Background(), 1*time.Minute) - defer cancel() - run, err := s.SdkClient().ExecuteWorkflow(ctx, sdkclient.StartWorkflowOptions{ + run, err := env.SdkClient().ExecuteWorkflow(env.Context(), sdkclient.StartWorkflowOptions{ ID: tvParent.WorkflowID(), TaskQueue: tvParent.TaskQueue().GetName(), }, "parent") s.NoError(err) // Set v2 as the current version of the deployment - s.setCurrentDeployment(tv2) + s.setCurrentDeployment(env, tv2) // Unblock parent to start the child - s.NoError(s.SdkClient().SignalWorkflow(ctx, run.GetID(), run.GetRunID(), "startChild", nil)) + s.NoError(env.SdkClient().SignalWorkflow(env.Context(), run.GetID(), run.GetRunID(), "startChild", nil)) // Verify that the child started and completed on the v2 worker. - s.Eventually(func() bool { - desc, err := s.SdkClient().DescribeWorkflowExecution(ctx, tvChild.WorkflowID(), "") - if err != nil { - return false - } + s.EventuallyWithT(func(t *assert.CollectT) { + desc, err := env.SdkClient().DescribeWorkflowExecution(env.Context(), tvChild.WorkflowID(), "") + require.NoError(t, err) execInfo := desc.GetWorkflowExecutionInfo() // Verify that the child workflow started and completed on the v1 worker - if execInfo.GetVersioningInfo().GetDeploymentVersion().GetBuildId() != tv2.BuildID() { - return false - } - if execInfo.GetStatus() != enumspb.WORKFLOW_EXECUTION_STATUS_COMPLETED { - return false - } - return true + require.Equal(t, tv2.BuildID(), execInfo.GetVersioningInfo().GetDeploymentVersion().GetBuildId()) + require.Equal(t, enumspb.WORKFLOW_EXECUTION_STATUS_COMPLETED, execInfo.GetStatus()) }, 10*time.Second, 100*time.Millisecond) // Verify that the parent workflow completed successfully. This shall only be possible if the child workflow // started and completed on the v2 worker. var result string - s.NoError(run.Get(ctx, &result)) + s.NoError(run.Get(env.Context(), &result)) s.Equal("v2", result) } func (s *Versioning3Suite) TestVersionedPoller_FailsWithEmptyNormalName() { - ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) - defer cancel() + env := s.setupEnv() - tv := testvars.New(s) + tv := env.Tv() // Simulate an old SDK polling a sticky task queue without providing normalName stickyTaskQueueWithoutNormalName := &taskqueuepb.TaskQueue{ @@ -4906,10 +4848,10 @@ func (s *Versioning3Suite) TestVersionedPoller_FailsWithEmptyNormalName() { } // Poll the sticky task queue - wfResponse, err := s.FrontendClient().PollWorkflowTaskQueue(ctx, &workflowservice.PollWorkflowTaskQueueRequest{ + wfResponse, err := env.FrontendClient().PollWorkflowTaskQueue(env.Context(), &workflowservice.PollWorkflowTaskQueueRequest{ TaskQueue: stickyTaskQueueWithoutNormalName, DeploymentOptions: tv.WorkerDeploymentOptions(true), - Namespace: s.Namespace().String(), + Namespace: env.Namespace().String(), Identity: tv.WorkerIdentity(), }) @@ -4919,10 +4861,10 @@ func (s *Versioning3Suite) TestVersionedPoller_FailsWithEmptyNormalName() { s.ErrorContains(err, "NormalName must be set on sticky queue when UseVersioning is true or DeploymentOptions are set.") // Poll activity task queue with an empty normalName (not possible, but conservative programming to safeguard against any potential bugs) - activityResponse, err := s.FrontendClient().PollActivityTaskQueue(ctx, &workflowservice.PollActivityTaskQueueRequest{ + activityResponse, err := env.FrontendClient().PollActivityTaskQueue(env.Context(), &workflowservice.PollActivityTaskQueueRequest{ TaskQueue: stickyTaskQueueWithoutNormalName, DeploymentOptions: tv.WorkerDeploymentOptions(true), - Namespace: s.Namespace().String(), + Namespace: env.Namespace().String(), Identity: tv.WorkerIdentity(), }) @@ -4933,10 +4875,12 @@ func (s *Versioning3Suite) TestVersionedPoller_FailsWithEmptyNormalName() { } func (s *Versioning3Suite) TestChildStartsWithParentRevision_SameTQ_TQLags() { - s.OverrideDynamicConfig(dynamicconfig.MatchingNumTaskqueueReadPartitions, 1) - s.OverrideDynamicConfig(dynamicconfig.MatchingNumTaskqueueWritePartitions, 1) + env := s.setupEnv( + testcore.WithDynamicConfig(dynamicconfig.MatchingNumTaskqueueReadPartitions, 1), + testcore.WithDynamicConfig(dynamicconfig.MatchingNumTaskqueueWritePartitions, 1), + ) - tvParent := testvars.New(s).WithBuildIDNumber(1).WithWorkflowIDNumber(1) + tvParent := env.Tv().WithBuildIDNumber(1).WithWorkflowIDNumber(1) tvChild := tvParent.WithWorkflowIDNumber(2) // Parent workflow waits for a signal before starting the child. @@ -4959,7 +4903,7 @@ func (s *Versioning3Suite) TestChildStartsWithParentRevision_SameTQ_TQLags() { } // Workers on parent and child TQs. Note: Both the parent and child workflows run on the same TQ. - w := worker.New(s.SdkClient(), tvParent.TaskQueue().GetName(), worker.Options{ + w := worker.New(env.SdkClient(), tvParent.TaskQueue().GetName(), worker.Options{ DeploymentOptions: worker.DeploymentOptions{Version: tvParent.SDKDeploymentVersion(), UseVersioning: true, DefaultVersioningBehavior: workflow.VersioningBehaviorAutoUpgrade}, }) w.RegisterWorkflowWithOptions(parent, workflow.RegisterOptions{Name: "parent", VersioningBehavior: workflow.VersioningBehaviorAutoUpgrade}) @@ -4968,60 +4912,49 @@ func (s *Versioning3Suite) TestChildStartsWithParentRevision_SameTQ_TQLags() { defer w.Stop() // Set v1 as the current version for the deployment - s.setCurrentDeployment(tvParent) + s.setCurrentDeployment(env, tvParent) // Start parent - ctx, cancel := context.WithTimeout(context.Background(), 1*time.Minute) - defer cancel() - run, err := s.SdkClient().ExecuteWorkflow(ctx, sdkclient.StartWorkflowOptions{ + run, err := env.SdkClient().ExecuteWorkflow(env.Context(), sdkclient.StartWorkflowOptions{ ID: tvParent.WorkflowID(), TaskQueue: tvParent.TaskQueue().GetName(), }, "parent") s.NoError(err) // Verify that the parent workflow has started on the v1 worker - s.Eventually(func() bool { - desc, err := s.SdkClient().DescribeWorkflowExecution(ctx, tvParent.WorkflowID(), "") - if err != nil { - return false - } - return desc.GetWorkflowExecutionInfo().GetVersioningInfo().GetDeploymentVersion().GetBuildId() == tvParent.BuildID() + s.EventuallyWithT(func(t *assert.CollectT) { + desc, err := env.SdkClient().DescribeWorkflowExecution(env.Context(), tvParent.WorkflowID(), "") + require.NoError(t, err) + require.Equal(t, tvParent.BuildID(), desc.GetWorkflowExecutionInfo().GetVersioningInfo().GetDeploymentVersion().GetBuildId()) }, 10*time.Second, 100*time.Millisecond) // Roll back the child TQ routing-config revision to simulate Routing Config lag in matching partitions (set v0 as current with older revision) tv0Child := tvChild.WithBuildIDNumber(0) - s.rollbackTaskQueueToVersion(tv0Child) + s.rollbackTaskQueueToVersion(env, tv0Child) //nolint:testifylint - idlePollerCtx, idlePollerCancel := context.WithTimeout(context.Background(), 10*time.Second) + idlePollerCtx, idlePollerCancel := context.WithTimeout(env.Context(), 10*time.Second) defer idlePollerCancel() - go s.idlePollWorkflow(idlePollerCtx, tv0Child, true, 10*time.Second, "workflow should not go to the old deployment") + go s.idlePollWorkflow(env, idlePollerCtx, tv0Child, true, 10*time.Second, "workflow should not go to the old deployment") // Unblock parent to start the child - s.NoError(s.SdkClient().SignalWorkflow(ctx, run.GetID(), run.GetRunID(), "startChild", nil)) + s.NoError(env.SdkClient().SignalWorkflow(env.Context(), run.GetID(), run.GetRunID(), "startChild", nil)) // Verify that the child started and completed on the v1 worker. - s.Eventually(func() bool { - desc, err := s.SdkClient().DescribeWorkflowExecution(ctx, tvChild.WorkflowID(), "") - if err != nil { - return false - } + s.EventuallyWithT(func(t *assert.CollectT) { + desc, err := env.SdkClient().DescribeWorkflowExecution(env.Context(), tvChild.WorkflowID(), "") + require.NoError(t, err) execInfo := desc.GetWorkflowExecutionInfo() // Verify that the child workflow started and completed on the v1 worker - if execInfo.GetVersioningInfo().GetDeploymentVersion().GetBuildId() != tvChild.BuildID() { - return false - } - if execInfo.GetStatus() != enumspb.WORKFLOW_EXECUTION_STATUS_COMPLETED { - return false - } - return true + require.Equal(t, tvChild.BuildID(), execInfo.GetVersioningInfo().GetDeploymentVersion().GetBuildId()) + require.Equal(t, enumspb.WORKFLOW_EXECUTION_STATUS_COMPLETED, execInfo.GetStatus()) }, 10*time.Second, 100*time.Millisecond) // Verify that the parent workflow completed successfully. This shall only be possible if the child workflow // started and completed on the v1 worker. var result string - s.NoError(run.Get(ctx, &result)) + s.NoError(run.Get(env.Context(), &result)) s.Equal("v1", result) } @@ -5029,11 +4962,13 @@ func (s *Versioning3Suite) TestChildStartsWithParentRevision_SameTQ_TQLags() { // TestChildStartsWithNoInheritedAutoUpgradeInfo_CrossTQ demonstrates that a child workflow of an AutoUpgrade parent, not sharing // the same task queue, starts with no inherited auto upgrade info. func (s *Versioning3Suite) TestChildStartsWithNoInheritedAutoUpgradeInfo_CrossTQ() { - s.OverrideDynamicConfig(dynamicconfig.MatchingNumTaskqueueReadPartitions, 1) - s.OverrideDynamicConfig(dynamicconfig.MatchingNumTaskqueueWritePartitions, 1) + env := s.setupEnv( + testcore.WithDynamicConfig(dynamicconfig.MatchingNumTaskqueueReadPartitions, 1), + testcore.WithDynamicConfig(dynamicconfig.MatchingNumTaskqueueWritePartitions, 1), + ) // Use different task queues that are in different deployments for both parent and child - tvParent := testvars.New(s).WithDeploymentSeriesNumber(1).WithBuildIDNumber(1).WithWorkflowIDNumber(1).WithTaskQueue("parent-tq") + tvParent := env.Tv().WithDeploymentSeriesNumber(1).WithBuildIDNumber(1).WithWorkflowIDNumber(1).WithTaskQueue("parent-tq") tvChild := tvParent.WithDeploymentSeriesNumber(2).WithBuildIDNumber(2).WithWorkflowIDNumber(2).WithTaskQueue("child-tq") // Parent workflow waits for a signal before starting the child. @@ -5054,7 +4989,7 @@ func (s *Versioning3Suite) TestChildStartsWithNoInheritedAutoUpgradeInfo_CrossTQ child := func(ctx workflow.Context) (string, error) { return "v1", nil } // Workers on parent and child TQs - w := worker.New(s.SdkClient(), tvParent.TaskQueue().GetName(), worker.Options{ + w := worker.New(env.SdkClient(), tvParent.TaskQueue().GetName(), worker.Options{ DeploymentOptions: worker.DeploymentOptions{Version: tvParent.SDKDeploymentVersion(), UseVersioning: true, DefaultVersioningBehavior: workflow.VersioningBehaviorAutoUpgrade}, }) w.RegisterWorkflowWithOptions(parent, workflow.RegisterOptions{Name: "parent", VersioningBehavior: workflow.VersioningBehaviorAutoUpgrade}) @@ -5062,7 +4997,7 @@ func (s *Versioning3Suite) TestChildStartsWithNoInheritedAutoUpgradeInfo_CrossTQ defer w.Stop() // Start child worker on child TQ - wChild := worker.New(s.SdkClient(), tvChild.TaskQueue().GetName(), worker.Options{ + wChild := worker.New(env.SdkClient(), tvChild.TaskQueue().GetName(), worker.Options{ DeploymentOptions: worker.DeploymentOptions{Version: tvChild.SDKDeploymentVersion(), UseVersioning: true, DefaultVersioningBehavior: workflow.VersioningBehaviorAutoUpgrade}, }) wChild.RegisterWorkflowWithOptions(child, workflow.RegisterOptions{Name: "child", VersioningBehavior: workflow.VersioningBehaviorAutoUpgrade}) @@ -5070,54 +5005,46 @@ func (s *Versioning3Suite) TestChildStartsWithNoInheritedAutoUpgradeInfo_CrossTQ defer wChild.Stop() // Set tvParent as the current version for the parent deployment - s.setCurrentDeployment(tvParent) + s.setCurrentDeployment(env, tvParent) // Set tvChild as the current version for the child deployment - s.setCurrentDeployment(tvChild) + s.setCurrentDeployment(env, tvChild) // Start parent - ctx, cancel := context.WithTimeout(context.Background(), 1*time.Minute) - defer cancel() - run, err := s.SdkClient().ExecuteWorkflow(ctx, sdkclient.StartWorkflowOptions{ + run, err := env.SdkClient().ExecuteWorkflow(env.Context(), sdkclient.StartWorkflowOptions{ ID: tvParent.WorkflowID(), TaskQueue: tvParent.TaskQueue().GetName(), }, "parent") s.NoError(err) // Unblock parent to start the child - s.NoError(s.SdkClient().SignalWorkflow(ctx, run.GetID(), run.GetRunID(), "startChild", nil)) + s.NoError(env.SdkClient().SignalWorkflow(env.Context(), run.GetID(), run.GetRunID(), "startChild", nil)) // Verify that the child workflow started and completed on it's worker, which shall verify that no AutoUpgradeInheritedInfo was inherited. - s.Eventually(func() bool { - s.Eventually(func() bool { - desc, err := s.SdkClient().DescribeWorkflowExecution(ctx, tvChild.WorkflowID(), "") - if err != nil { - return false - } - execInfo := desc.GetWorkflowExecutionInfo() + s.EventuallyWithT(func(t *assert.CollectT) { + desc, err := env.SdkClient().DescribeWorkflowExecution(env.Context(), tvChild.WorkflowID(), "") + require.NoError(t, err) + execInfo := desc.GetWorkflowExecutionInfo() - // Verify that the new run starts on the v1 worker - if execInfo.GetVersioningInfo().GetDeploymentVersion().GetBuildId() != tvChild.BuildID() { - return false - } - return true - }, 10*time.Second, 100*time.Millisecond) - return true + // Verify that the new run starts on the v1 worker + require.Equal(t, tvChild.BuildID(), execInfo.GetVersioningInfo().GetDeploymentVersion().GetBuildId()) }, 10*time.Second, 100*time.Millisecond) // Verify that the parent workflow completed successfully. // This shall only be possible if the child workflow started on it's worker and completed. var result string - s.NoError(run.Get(ctx, &result)) + s.NoError(run.Get(env.Context(), &result)) s.Equal("v1", result) } // Tests testing continue-as-new of an AutoUpgrade workflow using revision number mechanics. func (s *Versioning3Suite) TestContinueAsNewOfAutoUpgradeWorkflow_RevisionNumberMechanics() { - s.OverrideDynamicConfig(dynamicconfig.MatchingNumTaskqueueReadPartitions, 1) - s.OverrideDynamicConfig(dynamicconfig.MatchingNumTaskqueueWritePartitions, 1) + env := s.setupEnv( + testcore.WithDynamicConfig(dynamicconfig.MatchingNumTaskqueueReadPartitions, 1), + testcore.WithDynamicConfig(dynamicconfig.MatchingNumTaskqueueWritePartitions, 1), + ) - tv1 := testvars.New(s).WithBuildIDNumber(1).WithWorkflowIDNumber(1) + tv1 := env.Tv().WithBuildIDNumber(1).WithWorkflowIDNumber(1) // Workflow that waits for a signal before continuing-as-new canWorkflow := func(ctx workflow.Context, attempt int) (string, error) { @@ -5129,7 +5056,7 @@ func (s *Versioning3Suite) TestContinueAsNewOfAutoUpgradeWorkflow_RevisionNumber } // Worker on v1 - w := worker.New(s.SdkClient(), tv1.TaskQueue().GetName(), worker.Options{ + w := worker.New(env.SdkClient(), tv1.TaskQueue().GetName(), worker.Options{ DeploymentOptions: worker.DeploymentOptions{Version: tv1.SDKDeploymentVersion(), UseVersioning: true, DefaultVersioningBehavior: workflow.VersioningBehaviorAutoUpgrade}, }) w.RegisterWorkflowWithOptions(canWorkflow, workflow.RegisterOptions{Name: "canWorkflow", VersioningBehavior: workflow.VersioningBehaviorAutoUpgrade}) @@ -5137,65 +5064,52 @@ func (s *Versioning3Suite) TestContinueAsNewOfAutoUpgradeWorkflow_RevisionNumber defer w.Stop() // Set v1 to be the current version for the deployment - s.setCurrentDeployment(tv1) + s.setCurrentDeployment(env, tv1) // Start workflow - ctx, cancel := context.WithTimeout(context.Background(), 1*time.Minute) - defer cancel() - run, err := s.SdkClient().ExecuteWorkflow(ctx, sdkclient.StartWorkflowOptions{ + run, err := env.SdkClient().ExecuteWorkflow(env.Context(), sdkclient.StartWorkflowOptions{ ID: tv1.WorkflowID(), TaskQueue: tv1.TaskQueue().GetName(), }, "canWorkflow", 0) s.NoError(err) // Ensure the workflow has started on the v1 worker - s.Eventually(func() bool { - desc, err := s.SdkClient().DescribeWorkflowExecution(ctx, tv1.WorkflowID(), "") - if err != nil { - return false - } - return desc.GetWorkflowExecutionInfo().GetVersioningInfo().GetDeploymentVersion().GetBuildId() == tv1.BuildID() + s.EventuallyWithT(func(t *assert.CollectT) { + desc, err := env.SdkClient().DescribeWorkflowExecution(env.Context(), tv1.WorkflowID(), "") + require.NoError(t, err) + require.Equal(t, tv1.BuildID(), desc.GetWorkflowExecutionInfo().GetVersioningInfo().GetDeploymentVersion().GetBuildId()) }, 10*time.Second, 100*time.Millisecond) // Rollback the TaskQueueUserData to simulate task queue partition lag tv0 := tv1.WithBuildIDNumber(0) - s.rollbackTaskQueueToVersion(tv0) + s.rollbackTaskQueueToVersion(env, tv0) - idlePollerCtx, idlePollerCancel := context.WithTimeout(context.Background(), 10*time.Second) + idlePollerCtx, idlePollerCancel := context.WithTimeout(env.Context(), 10*time.Second) defer idlePollerCancel() - go s.idlePollWorkflow(idlePollerCtx, tv0, true, 10*time.Second, "workflow should not go to the old deployment") + go s.idlePollWorkflow(env, idlePollerCtx, tv0, true, 10*time.Second, "workflow should not go to the old deployment") // Signal the workflow to trigger CAN - s.NoError(s.SdkClient().SignalWorkflow(ctx, run.GetID(), run.GetRunID(), "triggerCAN", nil)) + s.NoError(env.SdkClient().SignalWorkflow(env.Context(), run.GetID(), run.GetRunID(), "triggerCAN", nil)) // Wait for the new run to start var newRunID string - s.Eventually(func() bool { - desc, err := s.SdkClient().DescribeWorkflowExecution(ctx, tv1.WorkflowID(), "") - if err != nil { - return false - } + s.EventuallyWithT(func(t *assert.CollectT) { + desc, err := env.SdkClient().DescribeWorkflowExecution(env.Context(), tv1.WorkflowID(), "") + require.NoError(t, err) execInfo := desc.GetWorkflowExecutionInfo() // Verify that the workflow truly CAN'ed and started a new run - if execInfo.GetStatus() == enumspb.WORKFLOW_EXECUTION_STATUS_COMPLETED { - newRunID = execInfo.GetExecution().GetRunId() - if newRunID == run.GetRunID() { - return false - } - } + require.Equal(t, enumspb.WORKFLOW_EXECUTION_STATUS_COMPLETED, execInfo.GetStatus()) + newRunID = execInfo.GetExecution().GetRunId() + require.NotEqual(t, run.GetRunID(), newRunID) // Verify that the new run starts on the v1 worker - if execInfo.GetVersioningInfo().GetDeploymentVersion().GetBuildId() != tv1.BuildID() { - return false - } - - return true + require.Equal(t, tv1.BuildID(), execInfo.GetVersioningInfo().GetDeploymentVersion().GetBuildId()) }, 10*time.Second, 100*time.Millisecond) // Verify that the workflow completed successfully on v1 var result string - s.NoError(run.Get(ctx, &result)) + s.NoError(run.Get(env.Context(), &result)) s.Equal("v1", result) } @@ -5204,11 +5118,8 @@ func (s *Versioning3Suite) TestContinueAsNewOfAutoUpgradeWorkflow_RevisionNumber // even if the task-queue partition's user data is rolled back to an older version. // If testContinueAsNew is true, tests a ContinueAsNew followed by retry; otherwise tests a direct retry of a workflow. // If testChildWorkflow is true, tests that a child workflow's retry doesn't bounce back (child spawned by parent with retry policy). -func (s *Versioning3Suite) testRetryNoBounceBack(testContinueAsNew bool, testChildWorkflow bool) { - s.OverrideDynamicConfig(dynamicconfig.MatchingNumTaskqueueReadPartitions, 1) - s.OverrideDynamicConfig(dynamicconfig.MatchingNumTaskqueueWritePartitions, 1) - - tv1 := testvars.New(s).WithBuildIDNumber(1) +func (s *Versioning3Suite) testRetryNoBounceBack(env *testcore.TestEnv, testContinueAsNew bool, testChildWorkflow bool) { + tv1 := env.Tv().WithBuildIDNumber(1) tv0 := tv1.WithBuildIDNumber(0) childWorkflowID := tv1.WorkflowID() + "-child" @@ -5292,11 +5203,8 @@ func (s *Versioning3Suite) testRetryNoBounceBack(testContinueAsNew bool, testChi } act := func() (string, error) { return "ok", nil } - ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute) - defer cancel() - // Worker for version 1 (current). - w1 := worker.New(s.SdkClient(), tv1.TaskQueue().GetName(), worker.Options{ + w1 := worker.New(env.SdkClient(), tv1.TaskQueue().GetName(), worker.Options{ DeploymentOptions: worker.DeploymentOptions{ Version: tv1.SDKDeploymentVersion(), UseVersioning: true, @@ -5314,15 +5222,15 @@ func (s *Versioning3Suite) testRetryNoBounceBack(testContinueAsNew bool, testChi defer w1.Stop() // Set v1 to be the current version for the deployment - s.setCurrentDeployment(tv1) + s.setCurrentDeployment(env, tv1) var wfID string var runIDBeforeRetry string if testChildWorkflow { // Start parent workflow (which spawns child with retry policy) - _, err := s.SdkClient().ExecuteWorkflow( - ctx, + _, err := env.SdkClient().ExecuteWorkflow( + env.Context(), sdkclient.StartWorkflowOptions{ TaskQueue: tv1.TaskQueue().GetName(), }, @@ -5331,20 +5239,18 @@ func (s *Versioning3Suite) testRetryNoBounceBack(testContinueAsNew bool, testChi s.NoError(err) // Wait for child workflow to be created and executing on v1 - s.Eventually(func() bool { - desc, err := s.SdkClient().DescribeWorkflowExecution(ctx, childWorkflowID, "") - if err != nil { - return false - } + s.EventuallyWithT(func(t *assert.CollectT) { + desc, err := env.SdkClient().DescribeWorkflowExecution(env.Context(), childWorkflowID, "") + require.NoError(t, err) runIDBeforeRetry = desc.GetWorkflowExecutionInfo().GetExecution().GetRunId() - return desc.GetWorkflowExecutionInfo().GetVersioningInfo().GetDeploymentVersion().GetBuildId() == tv1.BuildID() + require.Equal(t, tv1.BuildID(), desc.GetWorkflowExecutionInfo().GetVersioningInfo().GetDeploymentVersion().GetBuildId()) }, 10*time.Second, 100*time.Millisecond) wfID = childWorkflowID } else { // Start workflow with retry policy (so failure will trigger automatic retry). - run0, err := s.SdkClient().ExecuteWorkflow( - ctx, + run0, err := env.SdkClient().ExecuteWorkflow( + env.Context(), sdkclient.StartWorkflowOptions{ TaskQueue: tv1.TaskQueue().GetName(), RetryPolicy: &temporal.RetryPolicy{ @@ -5360,114 +5266,112 @@ func (s *Versioning3Suite) testRetryNoBounceBack(testContinueAsNew bool, testChi if testContinueAsNew { // Wait for ContinueAsNew to happen - s.Eventually(func() bool { - desc, err := s.SdkClient().DescribeWorkflowExecution(ctx, wfID, "") - s.NoError(err) - if err != nil { - return false - } + s.EventuallyWithT(func(t *assert.CollectT) { + desc, err := env.SdkClient().DescribeWorkflowExecution(env.Context(), wfID, "") + require.NoError(t, err) // After CAN, the run ID changes and we should see execution on v1 - if desc.GetWorkflowExecutionInfo().GetExecution().GetRunId() != run0.GetRunID() { - runIDBeforeRetry = desc.GetWorkflowExecutionInfo().GetExecution().GetRunId() - return desc.GetWorkflowExecutionInfo().GetVersioningInfo().GetDeploymentVersion().GetBuildId() == tv1.BuildID() - } - return false + runIDBeforeRetry = desc.GetWorkflowExecutionInfo().GetExecution().GetRunId() + require.NotEqual(t, run0.GetRunID(), runIDBeforeRetry) + require.Equal(t, tv1.BuildID(), desc.GetWorkflowExecutionInfo().GetVersioningInfo().GetDeploymentVersion().GetBuildId()) }, 10*time.Second, 100*time.Millisecond) } else { runIDBeforeRetry = run0.GetRunID() // Ensure initial run is executing on v1. - s.Eventually(func() bool { - desc, err := s.SdkClient().DescribeWorkflowExecution(ctx, wfID, runIDBeforeRetry) - s.NoError(err) - if err != nil { - return false - } - return desc.GetWorkflowExecutionInfo().GetVersioningInfo().GetDeploymentVersion().GetBuildId() == tv1.BuildID() + s.EventuallyWithT(func(t *assert.CollectT) { + desc, err := env.SdkClient().DescribeWorkflowExecution(env.Context(), wfID, runIDBeforeRetry) + require.NoError(t, err) + require.Equal(t, tv1.BuildID(), desc.GetWorkflowExecutionInfo().GetVersioningInfo().GetDeploymentVersion().GetBuildId()) }, 10*time.Second, 100*time.Millisecond) } } // Roll back the child TQ routing-config revision to simulate Routing Config lag in matching partitions (set v0 as current with older revision) - s.rollbackTaskQueueToVersion(tv0) + s.rollbackTaskQueueToVersion(env, tv0) // Start v0 pollers and ensure they don't receive a task - idlePollerCtx, idlePollerCancel := context.WithTimeout(context.Background(), 10*time.Second) + idlePollerCtx, idlePollerCancel := context.WithTimeout(env.Context(), 10*time.Second) defer idlePollerCancel() - go s.idlePollWorkflow(idlePollerCtx, tv0, true, 10*time.Second, "v0 poller should not receive a task") + go s.idlePollWorkflow(env, idlePollerCtx, tv0, true, 10*time.Second, "v0 poller should not receive a task") // Verify that the rollback propagated to all partitions - s.Eventually(func() bool { - ms, err := s.GetTestCluster().MatchingClient().GetTaskQueueUserData(context.Background(), &matchingservice.GetTaskQueueUserDataRequest{ - NamespaceId: s.NamespaceID().String(), + s.EventuallyWithT(func(t *assert.CollectT) { + ms, err := env.GetTestCluster().MatchingClient().GetTaskQueueUserData(env.Context(), &matchingservice.GetTaskQueueUserDataRequest{ + NamespaceId: env.NamespaceID().String(), TaskQueue: tv0.TaskQueue().GetName(), TaskQueueType: tqTypeWf, }) - s.NoError(err) + require.NoError(t, err) current, currentRevisionNumber, _, _, _, _, _, _ := worker_versioning.CalculateTaskQueueVersioningInfo(ms.GetUserData().GetData().GetPerType()[int32(tqTypeWf)].GetDeploymentData()) - return current.GetBuildId() == tv0.DeploymentVersion().GetBuildId() && currentRevisionNumber == 0 + require.Equal(t, tv0.DeploymentVersion().GetBuildId(), current.GetBuildId()) + require.Equal(t, int64(0), currentRevisionNumber) }, 10*time.Second, 100*time.Millisecond) // Trigger failure of the run to cause retry. - s.NoError(s.SdkClient().SignalWorkflow(ctx, wfID, runIDBeforeRetry, "proceed", nil)) + s.NoError(env.SdkClient().SignalWorkflow(env.Context(), wfID, runIDBeforeRetry, "proceed", nil)) // Wait for run to fail. - s.Eventually(func() bool { - desc, err := s.SdkClient().DescribeWorkflow(ctx, wfID, runIDBeforeRetry) - s.NoError(err) - if err != nil { - return false - } - return desc.Status == enumspb.WORKFLOW_EXECUTION_STATUS_FAILED + s.EventuallyWithT(func(t *assert.CollectT) { + desc, err := env.SdkClient().DescribeWorkflow(env.Context(), wfID, runIDBeforeRetry) + require.NoError(t, err) + require.Equal(t, enumspb.WORKFLOW_EXECUTION_STATUS_FAILED, desc.Status) }, 10*time.Second, 100*time.Millisecond) // Verify that retry run is still on v1 (didn't bounce back to v0) - s.Eventually(func() bool { - desc, err := s.SdkClient().DescribeWorkflowExecution(ctx, wfID, "") - s.NoError(err) - if err != nil { - return false - } + s.EventuallyWithT(func(t *assert.CollectT) { + desc, err := env.SdkClient().DescribeWorkflowExecution(env.Context(), wfID, "") + require.NoError(t, err) // After retry, there should be a new run - if desc.GetWorkflowExecutionInfo().GetExecution().GetRunId() != runIDBeforeRetry { - // Verify workflow (parent or child) is still on v1 - return desc.GetWorkflowExecutionInfo().GetVersioningInfo().GetDeploymentVersion().GetBuildId() == tv1.BuildID() - } - return false + require.NotEqual(t, runIDBeforeRetry, desc.GetWorkflowExecutionInfo().GetExecution().GetRunId()) + // Verify workflow (parent or child) is still on v1 + require.Equal(t, tv1.BuildID(), desc.GetWorkflowExecutionInfo().GetVersioningInfo().GetDeploymentVersion().GetBuildId()) }, 10*time.Second, 100*time.Millisecond) } func (s *Versioning3Suite) TestWorkflowRetry_AutoUpgrade_NoBounceBack() { - s.testRetryNoBounceBack(false, false) + env := s.setupEnv( + testcore.WithDynamicConfig(dynamicconfig.MatchingNumTaskqueueReadPartitions, 1), + testcore.WithDynamicConfig(dynamicconfig.MatchingNumTaskqueueWritePartitions, 1), + ) + s.testRetryNoBounceBack(env, false, false) } func (s *Versioning3Suite) TestWorkflowRetry_AutoUpgrade_AfterCAN_NoBounceBack() { - s.testRetryNoBounceBack(true, false) + env := s.setupEnv( + testcore.WithDynamicConfig(dynamicconfig.MatchingNumTaskqueueReadPartitions, 1), + testcore.WithDynamicConfig(dynamicconfig.MatchingNumTaskqueueWritePartitions, 1), + ) + s.testRetryNoBounceBack(env, true, false) } func (s *Versioning3Suite) TestWorkflowRetry_AutoUpgrade_ChildNoBounceBack() { - s.testRetryNoBounceBack(false, true) + env := s.setupEnv( + testcore.WithDynamicConfig(dynamicconfig.MatchingNumTaskqueueReadPartitions, 1), + testcore.WithDynamicConfig(dynamicconfig.MatchingNumTaskqueueWritePartitions, 1), + ) + s.testRetryNoBounceBack(env, false, true) } // The following tests test out the CheckTaskQueueVersionMembership RPC. func (s *Versioning3Suite) TestCheckTaskQueueVersionMembership() { - tv1 := testvars.New(s).WithBuildIDNumber(1) + env := s.setupEnv() + tv1 := env.Tv().WithBuildIDNumber(1) // No version exists in the task queue's userData as of now - s.Eventually(func() bool { - resp, err := s.GetTestCluster().MatchingClient().CheckTaskQueueVersionMembership(context.Background(), &matchingservice.CheckTaskQueueVersionMembershipRequest{ - NamespaceId: s.NamespaceID().String(), + s.EventuallyWithT(func(t *assert.CollectT) { + resp, err := env.GetTestCluster().MatchingClient().CheckTaskQueueVersionMembership(env.Context(), &matchingservice.CheckTaskQueueVersionMembershipRequest{ + NamespaceId: env.NamespaceID().String(), TaskQueue: tv1.TaskQueue().GetName(), TaskQueueType: tqTypeWf, Version: worker_versioning.DeploymentVersionFromDeployment(tv1.Deployment()), }) - s.NoError(err) - return !resp.GetIsMember() // the check should pass if no version is present + require.NoError(t, err) + require.False(t, resp.GetIsMember()) // the check should pass if no version is present }, 10*time.Second, 100*time.Millisecond) // Start v1 worker which shall register the version in the task queue - w1 := worker.New(s.SdkClient(), tv1.TaskQueue().GetName(), worker.Options{ + w1 := worker.New(env.SdkClient(), tv1.TaskQueue().GetName(), worker.Options{ DeploymentOptions: worker.DeploymentOptions{ Version: tv1.SDKDeploymentVersion(), UseVersioning: true, @@ -5477,25 +5381,25 @@ func (s *Versioning3Suite) TestCheckTaskQueueVersionMembership() { defer w1.Stop() // The version should eventually show up in the task queue's user data - s.Eventually(func() bool { - resp, err := s.GetTestCluster().MatchingClient().CheckTaskQueueVersionMembership(context.Background(), &matchingservice.CheckTaskQueueVersionMembershipRequest{ - NamespaceId: s.NamespaceID().String(), + s.EventuallyWithT(func(t *assert.CollectT) { + resp, err := env.GetTestCluster().MatchingClient().CheckTaskQueueVersionMembership(env.Context(), &matchingservice.CheckTaskQueueVersionMembershipRequest{ + NamespaceId: env.NamespaceID().String(), TaskQueue: tv1.TaskQueue().GetName(), TaskQueueType: tqTypeWf, Version: worker_versioning.DeploymentVersionFromDeployment(tv1.Deployment()), }) - s.NoError(err) - return resp.GetIsMember() + require.NoError(t, err) + require.True(t, resp.GetIsMember()) }, 10*time.Second, 100*time.Millisecond) } // validatePinnedVersionExistsInTaskQueue validates that the version, to be pinned, exists in the task queue. // TODO (future improvement): This can be further extended to validate the presence of any version instead of using the GetTaskQueueUserData RPC. -func (s *Versioning3Suite) validatePinnedVersionExistsInTaskQueue(tv *testvars.TestVars) { +func (s *Versioning3Suite) validatePinnedVersionExistsInTaskQueue(env *testcore.TestEnv, tv *testvars.TestVars) { s.EventuallyWithT(func(t *assert.CollectT) { a := require.New(t) - resp, err := s.GetTestCluster().MatchingClient().CheckTaskQueueVersionMembership(context.Background(), &matchingservice.CheckTaskQueueVersionMembershipRequest{ - NamespaceId: s.NamespaceID().String(), + resp, err := env.GetTestCluster().MatchingClient().CheckTaskQueueVersionMembership(env.Context(), &matchingservice.CheckTaskQueueVersionMembershipRequest{ + NamespaceId: env.NamespaceID().String(), TaskQueue: tv.TaskQueue().GetName(), TaskQueueType: tqTypeWf, Version: worker_versioning.DeploymentVersionFromDeployment(tv.Deployment()), @@ -5509,17 +5413,15 @@ func (s *Versioning3Suite) validatePinnedVersionExistsInTaskQueue(tv *testvars.T // TestMaxVersionsInTaskQueue tests that polling from a task queue with too many // versions returns a RESOURCE_EXHAUSTED_CAUSE_WORKER_DEPLOYMENT_LIMITS error func (s *Versioning3Suite) TestMaxVersionsInTaskQueue() { + maxVersions := 5 + env := s.setupEnv( + testcore.WithDynamicConfig(dynamicconfig.MatchingMaxVersionsInTaskQueue, maxVersions), + ) // This test is not dependent on a particular deployment workflow version, but we don't want it repetitively s.skipBeforeVersion(workerdeployment.VersionDataRevisionNumber) - ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) - defer cancel() - // Set a low limit for the test to make it faster - maxVersions := 5 - s.OverrideDynamicConfig(dynamicconfig.MatchingMaxVersionsInTaskQueue, maxVersions) - - tv := testvars.New(s) + tv := env.Tv() // Pre-populate the task queue with maxVersions different deployment versions // Each version will be in a separate deployment to ensure they count toward the limit @@ -5531,9 +5433,9 @@ func (s *Versioning3Suite) TestMaxVersionsInTaskQueue() { } deploymentName := tvVersion.DeploymentVersion().GetDeploymentName() - _, err := s.GetTestCluster().MatchingClient().SyncDeploymentUserData( - ctx, &matchingservice.SyncDeploymentUserDataRequest{ - NamespaceId: s.NamespaceID().String(), + _, err := env.GetTestCluster().MatchingClient().SyncDeploymentUserData( + env.Context(), &matchingservice.SyncDeploymentUserDataRequest{ + NamespaceId: env.NamespaceID().String(), TaskQueue: tv.TaskQueue().GetName(), TaskQueueTypes: []enumspb.TaskQueueType{tqTypeWf}, DeploymentName: deploymentName, @@ -5547,13 +5449,13 @@ func (s *Versioning3Suite) TestMaxVersionsInTaskQueue() { // This should fail with RESOURCE_EXHAUSTED_CAUSE_WORKER_DEPLOYMENT_LIMITS tvNewVersion := tv.WithDeploymentSeriesNumber(999).WithBuildIDNumber(999) - pollCtx, pollCancel := context.WithTimeout(ctx, 10*time.Second) + pollCtx, pollCancel := context.WithTimeout(env.Context(), 10*time.Second) defer pollCancel() - _, pollErr := s.FrontendClient().PollWorkflowTaskQueue(pollCtx, &workflowservice.PollWorkflowTaskQueueRequest{ + _, pollErr := env.FrontendClient().PollWorkflowTaskQueue(pollCtx, &workflowservice.PollWorkflowTaskQueueRequest{ TaskQueue: tvNewVersion.TaskQueue(), DeploymentOptions: tvNewVersion.WorkerDeploymentOptions(true), - Namespace: s.Namespace().String(), + Namespace: env.Namespace().String(), Identity: tvNewVersion.WorkerIdentity(), }) @@ -5570,15 +5472,16 @@ func (s *Versioning3Suite) TestMaxVersionsInTaskQueue() { } func (s *Versioning3Suite) TestActivityRetryAutoUpgradeDuringBackoff() { + env := s.setupEnv( + testcore.WithDynamicConfig(dynamicconfig.MatchingNumTaskqueueReadPartitions, 1), + testcore.WithDynamicConfig(dynamicconfig.MatchingNumTaskqueueWritePartitions, 1), + ) // This test verifies that when an activity is in retry backoff and the current deployment // changes, the workflow will auto-upgrade to the new deployment when the activity retry // executes, if the workflow's versioning behavior is AutoUpgrade. - s.OverrideDynamicConfig(dynamicconfig.MatchingNumTaskqueueReadPartitions, 1) - s.OverrideDynamicConfig(dynamicconfig.MatchingNumTaskqueueWritePartitions, 1) - numPollers := 4 - tv1 := testvars.New(s).WithBuildIDNumber(1) + tv1 := env.Tv().WithBuildIDNumber(1) tv2 := tv1.WithBuildIDNumber(2) // Track activity attempts and which version executed them @@ -5617,11 +5520,8 @@ func (s *Versioning3Suite) TestActivityRetryAutoUpgradeDuringBackoff() { return result, err } - ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute) - defer cancel() - // Create worker for v1 - w1 := worker.New(s.SdkClient(), tv1.TaskQueue().GetName(), worker.Options{ + w1 := worker.New(env.SdkClient(), tv1.TaskQueue().GetName(), worker.Options{ DeploymentOptions: worker.DeploymentOptions{ Version: tv1.SDKDeploymentVersion(), UseVersioning: true, @@ -5637,7 +5537,7 @@ func (s *Versioning3Suite) TestActivityRetryAutoUpgradeDuringBackoff() { defer w1.Stop() // Create worker for v2 - w2 := worker.New(s.SdkClient(), tv1.TaskQueue().GetName(), worker.Options{ + w2 := worker.New(env.SdkClient(), tv1.TaskQueue().GetName(), worker.Options{ DeploymentOptions: worker.DeploymentOptions{ Version: tv2.SDKDeploymentVersion(), UseVersioning: true, @@ -5653,12 +5553,12 @@ func (s *Versioning3Suite) TestActivityRetryAutoUpgradeDuringBackoff() { defer w2.Stop() // Set v1 as current deployment - s.setCurrentDeployment(tv1) - s.waitForDeploymentDataPropagation(tv1, versionStatusCurrent, false, tqTypeWf, tqTypeAct) + s.setCurrentDeployment(env, tv1) + s.waitForDeploymentDataPropagation(env, tv1, versionStatusCurrent, false, tqTypeWf, tqTypeAct) // Start the workflow - run, err := s.SdkClient().ExecuteWorkflow( - ctx, + run, err := env.SdkClient().ExecuteWorkflow( + env.Context(), sdkclient.StartWorkflowOptions{ TaskQueue: tv1.TaskQueue().GetName(), }, @@ -5667,27 +5567,25 @@ func (s *Versioning3Suite) TestActivityRetryAutoUpgradeDuringBackoff() { s.NoError(err) // Wait for the workflow to start on v1 - s.Eventually(func() bool { - desc, err := s.SdkClient().DescribeWorkflowExecution(ctx, run.GetID(), "") - if err != nil { - return false - } - return desc.GetWorkflowExecutionInfo().GetVersioningInfo().GetDeploymentVersion().GetBuildId() == tv1.BuildID() + s.EventuallyWithT(func(t *assert.CollectT) { + desc, err := env.SdkClient().DescribeWorkflowExecution(env.Context(), run.GetID(), "") + require.NoError(t, err) + require.Equal(t, tv1.BuildID(), desc.GetWorkflowExecutionInfo().GetVersioningInfo().GetDeploymentVersion().GetBuildId()) }, 10*time.Second, 100*time.Millisecond) // Wait for first activity attempt to fail (should be on v1) - s.Eventually(func() bool { - return v1AttemptCount.Load() == 1 + s.EventuallyWithT(func(t *assert.CollectT) { + require.Equal(t, int32(1), v1AttemptCount.Load()) }, 10*time.Second, 100*time.Millisecond) // Now the activity is in retry backoff. Change the current deployment to v2 // while the activity is waiting to retry. - s.setCurrentDeployment(tv2) - s.waitForDeploymentDataPropagation(tv2, versionStatusCurrent, false, tqTypeWf, tqTypeAct) + s.setCurrentDeployment(env, tv2) + s.waitForDeploymentDataPropagation(env, tv2, versionStatusCurrent, false, tqTypeWf, tqTypeAct) // Wait for the workflow to complete var result string - err = run.Get(ctx, &result) + err = run.Get(env.Context(), &result) s.NoError(err) // Verify results: @@ -5704,32 +5602,31 @@ func (s *Versioning3Suite) TestActivityRetryAutoUpgradeDuringBackoff() { s.Equal(int32(2), attemptCount.Load(), "Expected total of 2 activity attempts") // 5. Verify the workflow is now running on v2 - desc, err := s.SdkClient().DescribeWorkflowExecution(ctx, run.GetID(), "") + desc, err := env.SdkClient().DescribeWorkflowExecution(env.Context(), run.GetID(), "") s.NoError(err) s.Equal(tv2.BuildID(), desc.GetWorkflowExecutionInfo().GetVersioningInfo().GetDeploymentVersion().GetBuildId(), "Expected workflow to be on v2 after auto-upgrade") } func (s *Versioning3Suite) TestVersionedQueueUnload() { + env := s.setupEnv( + testcore.WithDynamicConfig(dynamicconfig.MatchingMaxTaskQueueIdleTime, 5*time.Second), + testcore.WithDynamicConfig(dynamicconfig.MatchingNumTaskqueueReadPartitions, 1), + testcore.WithDynamicConfig(dynamicconfig.MatchingNumTaskqueueWritePartitions, 1), + ) // Set MatchingMaxTaskQueueIdleTime to 5 seconds to shorten test duration - s.OverrideDynamicConfig(dynamicconfig.MatchingMaxTaskQueueIdleTime, 5*time.Second) - s.OverrideDynamicConfig(dynamicconfig.MatchingNumTaskqueueReadPartitions, 1) - s.OverrideDynamicConfig(dynamicconfig.MatchingNumTaskqueueWritePartitions, 1) - - ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) - defer cancel() - tv1 := testvars.New(s) + tv1 := env.Tv() - s.pollUntilRegistered(ctx, tv1) - s.setCurrentDeployment(tv1) - s.waitForDeploymentDataPropagation(tv1, versionStatusCurrent, false, tqTypeWf) + s.pollUntilRegistered(env, tv1) + s.setCurrentDeployment(env, tv1) + s.waitForDeploymentDataPropagation(env, tv1, versionStatusCurrent, false, tqTypeWf) // Start workflow with autoUpgrade behavior without any poller - s.startWorkflow(tv1, nil) + s.startWorkflow(env, tv1, nil) // Ensure the task is backlogged by checking stats - s.validateBacklogCount(tv1, tqTypeWf, 1) + s.validateBacklogCount(env, tv1, tqTypeWf, 1) // Ensure the default queue is loaded by sending long polls for user data for 10 seconds keepAliveDone := make(chan struct{}) @@ -5743,27 +5640,27 @@ func (s *Versioning3Suite) TestVersionedQueueUnload() { case <-timeout: return case <-ticker.C: - smallCtx, cancel := context.WithTimeout(ctx, 200*time.Millisecond) - _, _ = s.GetTestCluster().MatchingClient().GetTaskQueueUserData(smallCtx, &matchingservice.GetTaskQueueUserDataRequest{ - NamespaceId: s.NamespaceID().String(), + smallCtx, cancel := context.WithTimeout(env.Context(), 200*time.Millisecond) + _, _ = env.GetTestCluster().MatchingClient().GetTaskQueueUserData(smallCtx, &matchingservice.GetTaskQueueUserDataRequest{ + NamespaceId: env.NamespaceID().String(), OnlyIfLoaded: false, TaskQueue: tv1.TaskQueue().GetName(), TaskQueueType: tqTypeWf, WaitNewData: true, }) cancel() - case <-ctx.Done(): + case <-env.Context().Done(): return } } }() // Wait for 10 seconds - s.WaitForChannel(ctx, keepAliveDone) + env.WaitForChannel(env.Context(), keepAliveDone) // Send a poll from v1 and ensure the poller gets the task taskReceived := false - s.pollWftAndHandle(tv1, false, nil, + s.pollWftAndHandle(env, tv1, false, nil, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) taskReceived = true @@ -5773,42 +5670,44 @@ func (s *Versioning3Suite) TestVersionedQueueUnload() { s.True(taskReceived, "Poller should have received the task after keeping queue loaded") // Verify backlog is now empty - s.validateBacklogCount(tv1, tqTypeWf, 0) + s.validateBacklogCount(env, tv1, tqTypeWf, 0) } func (s *Versioning3Suite) TestTransitionDuringTransientTask_WithoutSignal() { - s.testTransitionDuringTransientTask(false) + env := s.setupEnv( + testcore.WithDynamicConfig(dynamicconfig.MatchingNumTaskqueueReadPartitions, 1), + testcore.WithDynamicConfig(dynamicconfig.MatchingNumTaskqueueWritePartitions, 1), + ) + s.testTransitionDuringTransientTask(env, false) } func (s *Versioning3Suite) TestTransitionDuringTransientTask_WithSignal() { - s.testTransitionDuringTransientTask(true) + env := s.setupEnv( + testcore.WithDynamicConfig(dynamicconfig.MatchingNumTaskqueueReadPartitions, 1), + testcore.WithDynamicConfig(dynamicconfig.MatchingNumTaskqueueWritePartitions, 1), + ) + s.testTransitionDuringTransientTask(env, true) } // TestTransitionDuringTransientTask verifies that a deployment version transition to v1 // is properly set when a workflow task fails and is retried, and that the transition completes // successfully after a signal is sent during the retry backoff period. -func (s *Versioning3Suite) testTransitionDuringTransientTask(withSignal bool) { - s.OverrideDynamicConfig(dynamicconfig.MatchingNumTaskqueueReadPartitions, 1) - s.OverrideDynamicConfig(dynamicconfig.MatchingNumTaskqueueWritePartitions, 1) - - ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second) - defer cancel() - +func (s *Versioning3Suite) testTransitionDuringTransientTask(env *testcore.TestEnv, withSignal bool) { // Create test vars for v1 - tv1 := testvars.New(s).WithBuildIDNumber(1) + tv1 := env.Tv().WithBuildIDNumber(1) - s.pollUntilRegistered(ctx, tv1, tqTypeWf, tqTypeAct) - s.waitForDeploymentDataPropagation(tv1, versionStatusInactive, false, tqTypeWf, tqTypeAct) + s.pollUntilRegistered(env, tv1, tqTypeWf, tqTypeAct) + s.waitForDeploymentDataPropagation(env, tv1, versionStatusInactive, false, tqTypeWf, tqTypeAct) // Start the workflow execution := &commonpb.WorkflowExecution{ WorkflowId: tv1.WorkflowID(), } - runID := s.startWorkflow(tv1, nil) + runID := s.startWorkflow(env, tv1, nil) execution.RunId = runID // Poll first workflow task and schedule two activities - s.unversionedPollWftAndHandle(tv1, false, nil, + s.unversionedPollWftAndHandle(env, tv1, false, nil, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) // Schedule two activities @@ -5816,14 +5715,14 @@ func (s *Versioning3Suite) testTransitionDuringTransientTask(withSignal bool) { }) // Complete the first activity to generate a new wft - s.unversionedPollActivityAndHandle(tv1, nil, + s.unversionedPollActivityAndHandle(env, tv1, nil, func(task *workflowservice.PollActivityTaskQueueResponse) (*workflowservice.RespondActivityTaskCompletedRequest, error) { s.NotNil(task) return respondActivity(), nil }) // Poll and fail the next workflow task - this will cause a transient WFT to be scheduled with retry backoff - s.unversionedPollWftAndHandle(tv1, false, nil, + s.unversionedPollWftAndHandle(env, tv1, false, nil, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) // Return error to fail the task @@ -5831,14 +5730,14 @@ func (s *Versioning3Suite) testTransitionDuringTransientTask(withSignal bool) { }) // Set v1 as current deployment - s.setCurrentDeployment(tv1) - s.waitForDeploymentDataPropagation(tv1, versionStatusCurrent, false, tqTypeWf, tqTypeAct) + s.setCurrentDeployment(env, tv1) + s.waitForDeploymentDataPropagation(env, tv1, versionStatusCurrent, false, tqTypeWf, tqTypeAct) if withSignal { // While during retry backoff, send a signal to the workflow // Signal would convert the transient task to normal because of the new history event being inserted after the first task failure - _, err := s.FrontendClient().SignalWorkflowExecution(ctx, &workflowservice.SignalWorkflowExecutionRequest{ - Namespace: s.Namespace().String(), + _, err := env.FrontendClient().SignalWorkflowExecution(env.Context(), &workflowservice.SignalWorkflowExecutionRequest{ + Namespace: env.Namespace().String(), WorkflowExecution: execution, SignalName: "test-signal", Input: tv1.Any().Payloads(), @@ -5848,12 +5747,12 @@ func (s *Versioning3Suite) testTransitionDuringTransientTask(withSignal bool) { } // Poll the second activity to cause transition to v1. - s.idlePollActivity(ctx, tv1, true, ver3MinPollTime, "should not get the activity because it started a transition") - s.verifyWorkflowVersioning(s.Assertions, tv1, vbUnspecified, nil, nil, tv1.DeploymentVersionTransition()) + s.idlePollActivity(env, tv1, true, ver3MinPollTime, "should not get the activity because it started a transition") + s.verifyWorkflowVersioning(env, s.Assertions, tv1, vbUnspecified, nil, nil, tv1.DeploymentVersionTransition()) // Print workflow describe and history - descResp, err := s.FrontendClient().DescribeWorkflowExecution(ctx, &workflowservice.DescribeWorkflowExecutionRequest{ - Namespace: s.Namespace().String(), + descResp, err := env.FrontendClient().DescribeWorkflowExecution(env.Context(), &workflowservice.DescribeWorkflowExecutionRequest{ + Namespace: env.Namespace().String(), Execution: execution, }) if err == nil && descResp != nil { @@ -5862,8 +5761,8 @@ func (s *Versioning3Suite) testTransitionDuringTransientTask(withSignal bool) { fmt.Println("=========================") } - histResp, err := s.FrontendClient().GetWorkflowExecutionHistory(ctx, &workflowservice.GetWorkflowExecutionHistoryRequest{ - Namespace: s.Namespace().String(), + histResp, err := env.FrontendClient().GetWorkflowExecutionHistory(env.Context(), &workflowservice.GetWorkflowExecutionHistoryRequest{ + Namespace: env.Namespace().String(), Execution: execution, }) if err == nil && histResp != nil { @@ -5876,11 +5775,11 @@ func (s *Versioning3Suite) testTransitionDuringTransientTask(withSignal bool) { // After the retry backoff, poll and complete the workflow task // This should complete the transition and the workflow should be on v1 - s.pollWftAndHandle(tv1, false, nil, + s.pollWftAndHandle(env, tv1, false, nil, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) // Verify that a deployment version transition to v1 is set in the workflow - s.verifyWorkflowVersioning(s.Assertions, tv1, vbUnspecified, nil, nil, tv1.DeploymentVersionTransition()) + s.verifyWorkflowVersioning(env, s.Assertions, tv1, vbUnspecified, nil, nil, tv1.DeploymentVersionTransition()) if withSignal { s.EqualHistory(` @@ -5924,45 +5823,44 @@ func (s *Versioning3Suite) testTransitionDuringTransientTask(withSignal bool) { }) // Verify the transition is completed and the workflow is on v1 - s.verifyWorkflowVersioning(s.Assertions, tv1, vbUnpinned, tv1.Deployment(), nil, nil) + s.verifyWorkflowVersioning(env, s.Assertions, tv1, vbUnpinned, tv1.Deployment(), nil, nil) } // TestPinnedCaN_NoAUOnCaN_NoInfiniteLoop tests that a pinned workflow that CAN's // without AUTO_UPGRADE as the InitialVersioningBehavior does not get stuck in an infinite CAN loop. func (s *Versioning3Suite) TestPinnedCaN_NoAUOnCaN_NoInfiniteLoop() { - ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second) - defer cancel() + env := s.setupEnv() - tv1 := testvars.New(s).WithBuildIDNumber(1) + tv1 := env.Tv().WithBuildIDNumber(1) tv2 := tv1.WithBuildIDNumber(2) // Start async poller for v1 that will handle the first WFT and declare pinned behavior wftCompleted := make(chan struct{}) - s.pollWftAndHandle(tv1, false, wftCompleted, + s.pollWftAndHandle(env, tv1, false, wftCompleted, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) return respondEmptyWft(tv1, false, vbPinned), nil }) // Wait for v1 to be registered, then set it as current - s.waitForDeploymentDataPropagation(tv1, versionStatusInactive, false, tqTypeWf) - s.setCurrentDeployment(tv1) + s.waitForDeploymentDataPropagation(env, tv1, versionStatusInactive, false, tqTypeWf) + s.setCurrentDeployment(env, tv1) // Start workflow — first WFT is handled by the async poller above (pinned on v1) - runID := s.startWorkflow(tv1, nil) + runID := s.startWorkflow(env, tv1, nil) execution := tv1.WithRunID(runID).WorkflowExecution() - s.WaitForChannel(ctx, wftCompleted) - s.verifyWorkflowVersioning(s.Assertions, tv1, vbPinned, tv1.Deployment(), nil, nil) + env.WaitForChannel(env.Context(), wftCompleted) + s.verifyWorkflowVersioning(env, s.Assertions, tv1, vbPinned, tv1.Deployment(), nil, nil) - s.pollUntilRegistered(ctx, tv2) - s.setCurrentDeployment(tv2) + s.pollUntilRegistered(env, tv2) + s.setCurrentDeployment(env, tv2) // Trigger WFT via signal - s.triggerNormalWFT(ctx, tv1, execution) + s.triggerNormalWFT(env, tv1, execution) // Process the WFT: verify targetWorkerDeploymentVersionChanged=true, then continue-as-new without AUTO_UPGRADE as the InitialVersioningBehavior // so that the new run starts on v1 (pinned). - s.pollWftAndHandle(tv1, false, nil, + s.pollWftAndHandle(env, tv1, false, nil, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) @@ -5997,7 +5895,7 @@ func (s *Versioning3Suite) TestPinnedCaN_NoAUOnCaN_NoInfiniteLoop() { }) // New run starts on v1 (pinned). Verify targetVersionChanged=false on first WFT — loop is broken. - s.pollWftAndHandle(tv1, false, nil, + s.pollWftAndHandle(env, tv1, false, nil, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) s.NotEqual(execution.RunId, task.WorkflowExecution.RunId, @@ -6011,55 +5909,54 @@ func (s *Versioning3Suite) TestPinnedCaN_NoAUOnCaN_NoInfiniteLoop() { return respondCompleteWorkflow(tv1, vbPinned), nil }) - s.verifyWorkflowVersioning(s.Assertions, tv1, vbPinned, tv1.Deployment(), nil, nil) + s.verifyWorkflowVersioning(env, s.Assertions, tv1, vbPinned, tv1.Deployment(), nil, nil) } // TestOverride_SuppressesTargetVersionChangedSignal tests that a versioning override // suppresses the target_worker_deployment_version_changed signal. When an operator // explicitly pins a workflow via override, routing changes should be irrelevant. func (s *Versioning3Suite) TestOverride_SuppressesTargetVersionChangedSignal() { - ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second) - defer cancel() + env := s.setupEnv() - tv1 := testvars.New(s).WithBuildIDNumber(1) + tv1 := env.Tv().WithBuildIDNumber(1) tv2 := tv1.WithBuildIDNumber(2) // Start async poller for v1 that will handle the first WFT and declare pinned behavior wftCompleted := make(chan struct{}) - s.pollWftAndHandle(tv1, false, wftCompleted, + s.pollWftAndHandle(env, tv1, false, wftCompleted, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) return respondEmptyWft(tv1, false, vbPinned), nil }) // Wait for v1 to be registered, then set it as current - s.waitForDeploymentDataPropagation(tv1, versionStatusInactive, false, tqTypeWf) - s.setCurrentDeployment(tv1) + s.waitForDeploymentDataPropagation(env, tv1, versionStatusInactive, false, tqTypeWf) + s.setCurrentDeployment(env, tv1) // Start workflow — first WFT is handled by the async poller above (pinned on v1) - s.startWorkflow(tv1, nil) + s.startWorkflow(env, tv1, nil) execution := tv1.WorkflowExecution() - s.WaitForChannel(ctx, wftCompleted) - s.verifyWorkflowVersioning(s.Assertions, tv1, vbPinned, tv1.Deployment(), nil, nil) + env.WaitForChannel(env.Context(), wftCompleted) + s.verifyWorkflowVersioning(env, s.Assertions, tv1, vbPinned, tv1.Deployment(), nil, nil) // Apply pinned override to v1 via UpdateWorkflowExecutionOptions override := s.makePinnedOverride(tv1) - _, err := s.FrontendClient().UpdateWorkflowExecutionOptions(ctx, &workflowservice.UpdateWorkflowExecutionOptionsRequest{ - Namespace: s.Namespace().String(), + _, err := env.FrontendClient().UpdateWorkflowExecutionOptions(env.Context(), &workflowservice.UpdateWorkflowExecutionOptionsRequest{ + Namespace: env.Namespace().String(), WorkflowExecution: execution, WorkflowExecutionOptions: &workflowpb.WorkflowExecutionOptions{VersioningOverride: override}, UpdateMask: &fieldmaskpb.FieldMask{Paths: []string{"versioning_override"}}, }) s.NoError(err) - s.pollUntilRegistered(ctx, tv2) - s.setCurrentDeployment(tv2) + s.pollUntilRegistered(env, tv2) + s.setCurrentDeployment(env, tv2) // Trigger WFT via signal - s.triggerNormalWFT(ctx, tv1, execution) + s.triggerNormalWFT(env, tv1, execution) // Verify signal=false: override suppresses the target version changed signal - s.pollWftAndHandle(tv1, false, nil, + s.pollWftAndHandle(env, tv1, false, nil, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) for _, event := range task.History.GetEvents() { @@ -6077,39 +5974,38 @@ func (s *Versioning3Suite) TestOverride_SuppressesTargetVersionChangedSignal() { // routing target changes. AutoUpgrade workflows will naturally transition to the // new version on the next WFT, so signaling CaN is unnecessary. func (s *Versioning3Suite) TestAutoUpgrade_SuppressesTargetVersionChangedSignal() { - ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second) - defer cancel() + env := s.setupEnv() - tv1 := testvars.New(s).WithBuildIDNumber(1) + tv1 := env.Tv().WithBuildIDNumber(1) tv2 := tv1.WithBuildIDNumber(2) // Start async poller for v1 that will handle the first WFT and declare AutoUpgrade behavior wftCompleted := make(chan struct{}) - s.pollWftAndHandle(tv1, false, wftCompleted, + s.pollWftAndHandle(env, tv1, false, wftCompleted, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) return respondEmptyWft(tv1, false, vbUnpinned), nil }) // Wait for v1 to be registered, then set it as current - s.waitForDeploymentDataPropagation(tv1, versionStatusInactive, false, tqTypeWf) - s.setCurrentDeployment(tv1) + s.waitForDeploymentDataPropagation(env, tv1, versionStatusInactive, false, tqTypeWf) + s.setCurrentDeployment(env, tv1) // Start workflow — first WFT is handled by the async poller above (AutoUpgrade on v1) - runID := s.startWorkflow(tv1, nil) + runID := s.startWorkflow(env, tv1, nil) execution := tv1.WithRunID(runID).WorkflowExecution() - s.WaitForChannel(ctx, wftCompleted) - s.verifyWorkflowVersioning(s.Assertions, tv1, vbUnpinned, tv1.Deployment(), nil, nil) + env.WaitForChannel(env.Context(), wftCompleted) + s.verifyWorkflowVersioning(env, s.Assertions, tv1, vbUnpinned, tv1.Deployment(), nil, nil) - s.pollUntilRegistered(ctx, tv2) - s.setCurrentDeployment(tv2) + s.pollUntilRegistered(env, tv2) + s.setCurrentDeployment(env, tv2) // Trigger WFT via signal - s.triggerNormalWFT(ctx, tv1, execution) + s.triggerNormalWFT(env, tv1, execution) // AutoUpgrade workflow transitions to v2, so the WFT is dispatched to v2 pollers. // Verify signal=false: AutoUpgrade suppresses the target version changed signal. - s.pollWftAndHandle(tv2, false, nil, + s.pollWftAndHandle(env, tv2, false, nil, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) for _, event := range task.History.GetEvents() { @@ -6128,39 +6024,38 @@ func (s *Versioning3Suite) TestAutoUpgrade_SuppressesTargetVersionChangedSignal( // This ensures condition 4 only suppresses the signal when the target is unchanged, // not when a genuinely new target appears. func (s *Versioning3Suite) TestPinnedCaN_TargetChangesAgain_SignalsTrue() { - ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second) - defer cancel() + env := s.setupEnv() - tv1 := testvars.New(s).WithBuildIDNumber(1) + tv1 := env.Tv().WithBuildIDNumber(1) tv2 := tv1.WithBuildIDNumber(2) tv3 := tv1.WithBuildIDNumber(3) // Start async poller for v1 that will handle the first WFT and declare pinned behavior wftCompleted := make(chan struct{}) - s.pollWftAndHandle(tv1, false, wftCompleted, + s.pollWftAndHandle(env, tv1, false, wftCompleted, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) return respondEmptyWft(tv1, false, vbPinned), nil }) // Wait for v1 to be registered, then set it as current - s.waitForDeploymentDataPropagation(tv1, versionStatusInactive, false, tqTypeWf) - s.setCurrentDeployment(tv1) + s.waitForDeploymentDataPropagation(env, tv1, versionStatusInactive, false, tqTypeWf) + s.setCurrentDeployment(env, tv1) // Start workflow — first WFT is handled by the async poller above (pinned on v1) - runID := s.startWorkflow(tv1, nil) + runID := s.startWorkflow(env, tv1, nil) execution := tv1.WithRunID(runID).WorkflowExecution() - s.WaitForChannel(ctx, wftCompleted) - s.verifyWorkflowVersioning(s.Assertions, tv1, vbPinned, tv1.Deployment(), nil, nil) + env.WaitForChannel(env.Context(), wftCompleted) + s.verifyWorkflowVersioning(env, s.Assertions, tv1, vbPinned, tv1.Deployment(), nil, nil) - s.pollUntilRegistered(ctx, tv2) - s.setCurrentDeployment(tv2) + s.pollUntilRegistered(env, tv2) + s.setCurrentDeployment(env, tv2) // Trigger WFT via signal - s.triggerNormalWFT(ctx, tv1, execution) + s.triggerNormalWFT(env, tv1, execution) // Process the WFT: verify targetWorkerDeploymentVersionChanged=true, then CaN without AU - s.pollWftAndHandle(tv1, false, nil, + s.pollWftAndHandle(env, tv1, false, nil, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) historyEvents := task.History.GetEvents() @@ -6194,13 +6089,13 @@ func (s *Versioning3Suite) TestPinnedCaN_TargetChangesAgain_SignalsTrue() { }) // Now change current to v3 (target changes from v2 to v3) - s.pollUntilRegistered(ctx, tv3) - s.setCurrentDeployment(tv3) - s.waitForDeploymentDataPropagation(tv3, versionStatusCurrent, false, tqTypeWf) + s.pollUntilRegistered(env, tv3) + s.setCurrentDeployment(env, tv3) + s.waitForDeploymentDataPropagation(env, tv3, versionStatusCurrent, false, tqTypeWf) // New run starts on v1 (pinned). Since target changed from v2→v3, // condition 4 should NOT suppress — targetVersionChanged should be true. - s.pollWftAndHandle(tv1, false, nil, + s.pollWftAndHandle(env, tv1, false, nil, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) s.NotEqual(execution.RunId, task.WorkflowExecution.RunId, @@ -6231,37 +6126,36 @@ func (s *Versioning3Suite) TestPinnedCaN_TargetChangesAgain_SignalsTrue() { // 4. Remove override — target is still v2 // 5. Signal → assert true (declined should have been cleared by override) func (s *Versioning3Suite) TestRemoveOverride_ClearsDeclinedState() { - ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second) - defer cancel() + env := s.setupEnv() - tv1 := testvars.New(s).WithBuildIDNumber(1) + tv1 := env.Tv().WithBuildIDNumber(1) tv2 := tv1.WithBuildIDNumber(2) // Start async poller for v1 that will handle the first WFT and declare pinned behavior wftCompleted := make(chan struct{}) - s.pollWftAndHandle(tv1, false, wftCompleted, + s.pollWftAndHandle(env, tv1, false, wftCompleted, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) return respondEmptyWft(tv1, false, vbPinned), nil }) // Wait for v1 to be registered, then set it as current - s.waitForDeploymentDataPropagation(tv1, versionStatusInactive, false, tqTypeWf) - s.setCurrentDeployment(tv1) + s.waitForDeploymentDataPropagation(env, tv1, versionStatusInactive, false, tqTypeWf) + s.setCurrentDeployment(env, tv1) // Start workflow — first WFT is handled by the async poller above (pinned on v1) - runID := s.startWorkflow(tv1, nil) + runID := s.startWorkflow(env, tv1, nil) execution := tv1.WithRunID(runID).WorkflowExecution() - s.WaitForChannel(ctx, wftCompleted) - s.verifyWorkflowVersioning(s.Assertions, tv1, vbPinned, tv1.Deployment(), nil, nil) + env.WaitForChannel(env.Context(), wftCompleted) + s.verifyWorkflowVersioning(env, s.Assertions, tv1, vbPinned, tv1.Deployment(), nil, nil) // Set v2 as current, trigger signal - s.pollUntilRegistered(ctx, tv2) - s.setCurrentDeployment(tv2) - s.triggerNormalWFT(ctx, tv1, execution) + s.pollUntilRegistered(env, tv2) + s.setCurrentDeployment(env, tv2) + s.triggerNormalWFT(env, tv1, execution) // Run 1: targetDeploymentVersionChanged=true → CaN without AU (decline v2) - s.pollWftAndHandle(tv1, false, nil, + s.pollWftAndHandle(env, tv1, false, nil, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) var lastStarted *historypb.HistoryEvent @@ -6293,7 +6187,7 @@ func (s *Versioning3Suite) TestRemoveOverride_ClearsDeclinedState() { }) // CaN run: declined=v2, target=v2 → case 4 suppresses - s.pollWftAndHandle(tv1, false, nil, + s.pollWftAndHandle(env, tv1, false, nil, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) s.NotEqual(execution.RunId, task.WorkflowExecution.RunId, @@ -6314,8 +6208,8 @@ func (s *Versioning3Suite) TestRemoveOverride_ClearsDeclinedState() { // Set override on the CaN run override := s.makePinnedOverride(tv1) - _, err := s.FrontendClient().UpdateWorkflowExecutionOptions(ctx, &workflowservice.UpdateWorkflowExecutionOptionsRequest{ - Namespace: s.Namespace().String(), + _, err := env.FrontendClient().UpdateWorkflowExecutionOptions(env.Context(), &workflowservice.UpdateWorkflowExecutionOptionsRequest{ + Namespace: env.Namespace().String(), WorkflowExecution: execution, WorkflowExecutionOptions: &workflowpb.WorkflowExecutionOptions{VersioningOverride: override}, UpdateMask: &fieldmaskpb.FieldMask{Paths: []string{"versioning_override"}}, @@ -6323,8 +6217,8 @@ func (s *Versioning3Suite) TestRemoveOverride_ClearsDeclinedState() { s.NoError(err) // Trigger a WFT while override is active so case 1 fires (and should clear declined) - s.triggerNormalWFT(ctx, tv1, execution) - s.pollWftAndHandle(tv1, false, nil, + s.triggerNormalWFT(env, tv1, execution) + s.pollWftAndHandle(env, tv1, false, nil, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) var lastStarted *historypb.HistoryEvent @@ -6340,8 +6234,8 @@ func (s *Versioning3Suite) TestRemoveOverride_ClearsDeclinedState() { }) // Remove the override - _, err = s.FrontendClient().UpdateWorkflowExecutionOptions(ctx, &workflowservice.UpdateWorkflowExecutionOptionsRequest{ - Namespace: s.Namespace().String(), + _, err = env.FrontendClient().UpdateWorkflowExecutionOptions(env.Context(), &workflowservice.UpdateWorkflowExecutionOptionsRequest{ + Namespace: env.Namespace().String(), WorkflowExecution: execution, WorkflowExecutionOptions: &workflowpb.WorkflowExecutionOptions{VersioningOverride: nil}, UpdateMask: &fieldmaskpb.FieldMask{Paths: []string{"versioning_override"}}, @@ -6349,9 +6243,9 @@ func (s *Versioning3Suite) TestRemoveOverride_ClearsDeclinedState() { s.NoError(err) // Trigger WFT — target is still v2, but declined should have been cleared - s.triggerNormalWFT(ctx, tv1, execution) + s.triggerNormalWFT(env, tv1, execution) - s.pollWftAndHandle(tv1, false, nil, + s.pollWftAndHandle(env, tv1, false, nil, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) var lastStarted *historypb.HistoryEvent @@ -6381,47 +6275,46 @@ func (s *Versioning3Suite) TestRemoveOverride_ClearsDeclinedState() { // 7. Set v4 as current (fresh version with higher revision) // 8. Trigger WFT → assert targetDeploymentVersionChanged=true (new revision > declined revision) func (s *Versioning3Suite) TestStalePartition_RevisionSuppressesTrampolining() { - s.OverrideDynamicConfig(dynamicconfig.MatchingNumTaskqueueReadPartitions, 1) - s.OverrideDynamicConfig(dynamicconfig.MatchingNumTaskqueueWritePartitions, 1) - - ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second) - defer cancel() + env := s.setupEnv( + testcore.WithDynamicConfig(dynamicconfig.MatchingNumTaskqueueReadPartitions, 1), + testcore.WithDynamicConfig(dynamicconfig.MatchingNumTaskqueueWritePartitions, 1), + ) - tv1 := testvars.New(s).WithBuildIDNumber(1) + tv1 := env.Tv().WithBuildIDNumber(1) tv2 := tv1.WithBuildIDNumber(2) tv3 := tv1.WithBuildIDNumber(3) // Start async poller for v1 that handles first WFT and declares pinned behavior wftCompleted := make(chan struct{}) - s.pollWftAndHandle(tv1, false, wftCompleted, + s.pollWftAndHandle(env, tv1, false, wftCompleted, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) return respondEmptyWft(tv1, false, vbPinned), nil }) // Register v1 and set it as current - s.waitForDeploymentDataPropagation(tv1, versionStatusInactive, false, tqTypeWf) - s.setCurrentDeployment(tv1) + s.waitForDeploymentDataPropagation(env, tv1, versionStatusInactive, false, tqTypeWf) + s.setCurrentDeployment(env, tv1) // Start workflow — first WFT handled by async poller (pinned on v1) - runID := s.startWorkflow(tv1, nil) + runID := s.startWorkflow(env, tv1, nil) execution := tv1.WithRunID(runID).WorkflowExecution() - s.WaitForChannel(ctx, wftCompleted) - s.verifyWorkflowVersioning(s.Assertions, tv1, vbPinned, tv1.Deployment(), nil, nil) + env.WaitForChannel(env.Context(), wftCompleted) + s.verifyWorkflowVersioning(env, s.Assertions, tv1, vbPinned, tv1.Deployment(), nil, nil) // Register v2, set v2 as current (revision increments) - s.idlePollWorkflow(ctx, tv2, true, ver3MinPollTime, "v2 poller registration") - s.setCurrentDeployment(tv2) + s.idlePollWorkflow(env, env.Context(), tv2, true, ver3MinPollTime, "v2 poller registration") + s.setCurrentDeployment(env, tv2) // Register v3, set v3 as current (revision increments again) - s.idlePollWorkflow(ctx, tv3, true, ver3MinPollTime, "v3 poller registration") - s.setCurrentDeployment(tv3) + s.idlePollWorkflow(env, env.Context(), tv3, true, ver3MinPollTime, "v3 poller registration") + s.setCurrentDeployment(env, tv3) // Trigger WFT — target should be v3 with a high revision - s.triggerNormalWFT(ctx, tv1, execution) + s.triggerNormalWFT(env, tv1, execution) // Process: targetDeploymentVersionChanged=true → CaN without AU (decline v3) - s.pollWftAndHandle(tv1, false, nil, + s.pollWftAndHandle(env, tv1, false, nil, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) var lastStarted *historypb.HistoryEvent @@ -6453,7 +6346,7 @@ func (s *Versioning3Suite) TestStalePartition_RevisionSuppressesTrampolining() { }) // CaN run: first WFT — declined=v3 propagated from previous run - s.pollWftAndHandle(tv1, false, nil, + s.pollWftAndHandle(env, tv1, false, nil, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) s.NotEqual(execution.RunId, task.WorkflowExecution.RunId, @@ -6463,13 +6356,13 @@ func (s *Versioning3Suite) TestStalePartition_RevisionSuppressesTrampolining() { }) // Simulate stale partition: roll back task queue to v2 with revision 0 - s.rollbackTaskQueueToVersion(tv2) + s.rollbackTaskQueueToVersion(env, tv2) // Trigger WFT with stale data — target is now v2 at revision 0 - s.triggerNormalWFT(ctx, tv1, execution) + s.triggerNormalWFT(env, tv1, execution) // Assert: revision 0 <= declined revision → signal suppressed - s.pollWftAndHandle(tv1, false, nil, + s.pollWftAndHandle(env, tv1, false, nil, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) var lastStarted *historypb.HistoryEvent @@ -6487,15 +6380,15 @@ func (s *Versioning3Suite) TestStalePartition_RevisionSuppressesTrampolining() { // Set a new v4 as current — this produces a revision strictly higher than // the declined revision, simulating an up-to-date partition with fresh data. tv4 := tv1.WithBuildIDNumber(4) - s.idlePollWorkflow(ctx, tv4, true, ver3MinPollTime, "v4 poller registration") - s.setCurrentDeployment(tv4) - s.waitForDeploymentDataPropagation(tv4, versionStatusCurrent, false, tqTypeWf) + s.idlePollWorkflow(env, env.Context(), tv4, true, ver3MinPollTime, "v4 poller registration") + s.setCurrentDeployment(env, tv4) + s.waitForDeploymentDataPropagation(env, tv4, versionStatusCurrent, false, tqTypeWf) // Trigger WFT with fresh data — target is v4 at higher revision - s.triggerNormalWFT(ctx, tv1, execution) + s.triggerNormalWFT(env, tv1, execution) // Assert: new revision > declined revision → signal fires - s.pollWftAndHandle(tv1, false, nil, + s.pollWftAndHandle(env, tv1, false, nil, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) var lastStarted *historypb.HistoryEvent @@ -6526,10 +6419,12 @@ func (s *Versioning3Suite) TestStalePartition_RevisionSuppressesTrampolining() { // 5. Retry run: OnStart=v2 (preserved from CaN decision), target=v2 // 6. Assert targetDeploymentVersionChanged=false (v2 == v2, decline is preserved) func (s *Versioning3Suite) TestRetryOfDeclinedCaN_SignalsOnNewTarget() { - s.OverrideDynamicConfig(dynamicconfig.MatchingNumTaskqueueReadPartitions, 1) - s.OverrideDynamicConfig(dynamicconfig.MatchingNumTaskqueueWritePartitions, 1) + env := s.setupEnv( + testcore.WithDynamicConfig(dynamicconfig.MatchingNumTaskqueueReadPartitions, 1), + testcore.WithDynamicConfig(dynamicconfig.MatchingNumTaskqueueWritePartitions, 1), + ) - tv1 := testvars.New(s).WithBuildIDNumber(1) + tv1 := env.Tv().WithBuildIDNumber(1) tv2 := tv1.WithBuildIDNumber(2) numPollers := 4 @@ -6552,10 +6447,7 @@ func (s *Versioning3Suite) TestRetryOfDeclinedCaN_SignalsOnNewTarget() { return "done", nil } - ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute) - defer cancel() - - w1 := worker.New(s.SdkClient(), tv1.TaskQueue().GetName(), worker.Options{ + w1 := worker.New(env.SdkClient(), tv1.TaskQueue().GetName(), worker.Options{ DeploymentOptions: worker.DeploymentOptions{ Version: tv1.SDKDeploymentVersion(), UseVersioning: true, @@ -6569,11 +6461,11 @@ func (s *Versioning3Suite) TestRetryOfDeclinedCaN_SignalsOnNewTarget() { defer w1.Stop() // Set v1 as current. - s.setCurrentDeployment(tv1) - s.waitForDeploymentDataPropagation(tv1, versionStatusCurrent, false, tqTypeWf) + s.setCurrentDeployment(env, tv1) + s.waitForDeploymentDataPropagation(env, tv1, versionStatusCurrent, false, tqTypeWf) // Start workflow with retry policy. - run0, err := s.SdkClient().ExecuteWorkflow(ctx, + run0, err := env.SdkClient().ExecuteWorkflow(env.Context(), sdkclient.StartWorkflowOptions{ TaskQueue: tv1.TaskQueue().GetName(), RetryPolicy: &temporal.RetryPolicy{ @@ -6586,59 +6478,51 @@ func (s *Versioning3Suite) TestRetryOfDeclinedCaN_SignalsOnNewTarget() { wfID := run0.GetID() // Wait for workflow to be running on v1. - s.Eventually(func() bool { - desc, err := s.SdkClient().DescribeWorkflowExecution(ctx, wfID, run0.GetRunID()) - if err != nil { - return false - } - return desc.GetWorkflowExecutionInfo().GetVersioningInfo().GetDeploymentVersion().GetBuildId() == tv1.BuildID() + s.EventuallyWithT(func(t *assert.CollectT) { + desc, err := env.SdkClient().DescribeWorkflowExecution(env.Context(), wfID, run0.GetRunID()) + require.NoError(t, err) + require.Equal(t, tv1.BuildID(), desc.GetWorkflowExecutionInfo().GetVersioningInfo().GetDeploymentVersion().GetBuildId()) }, 10*time.Second, 100*time.Millisecond) // Set v2 as current, signal workflow to CaN without AU (decline upgrade). - s.pollUntilRegistered(ctx, tv2) - s.setCurrentDeployment(tv2) - s.waitForDeploymentDataPropagation(tv2, versionStatusCurrent, false, tqTypeWf) - s.NoError(s.SdkClient().SignalWorkflow(ctx, wfID, run0.GetRunID(), "proceed", nil)) + s.pollUntilRegistered(env, tv2) + s.setCurrentDeployment(env, tv2) + s.waitForDeploymentDataPropagation(env, tv2, versionStatusCurrent, false, tqTypeWf) + s.NoError(env.SdkClient().SignalWorkflow(env.Context(), wfID, run0.GetRunID(), "proceed", nil)) // Wait for CaN to happen — new run on v1. var canRunID string - s.Eventually(func() bool { - desc, err := s.SdkClient().DescribeWorkflowExecution(ctx, wfID, "") - if err != nil { - return false - } + s.EventuallyWithT(func(t *assert.CollectT) { + desc, err := env.SdkClient().DescribeWorkflowExecution(env.Context(), wfID, "") + require.NoError(t, err) canRunID = desc.GetWorkflowExecutionInfo().GetExecution().GetRunId() - return canRunID != run0.GetRunID() && - desc.GetWorkflowExecutionInfo().GetVersioningInfo().GetDeploymentVersion().GetBuildId() == tv1.BuildID() + require.NotEqual(t, run0.GetRunID(), canRunID) + require.Equal(t, tv1.BuildID(), desc.GetWorkflowExecutionInfo().GetVersioningInfo().GetDeploymentVersion().GetBuildId()) }, 10*time.Second, 100*time.Millisecond) // Signal CaN run to fail (triggers server retry). Target remains v2. - s.NoError(s.SdkClient().SignalWorkflow(ctx, wfID, canRunID, "proceed", nil)) + s.NoError(env.SdkClient().SignalWorkflow(env.Context(), wfID, canRunID, "proceed", nil)) // Wait for CaN run to fail. - s.Eventually(func() bool { - desc, err := s.SdkClient().DescribeWorkflow(ctx, wfID, canRunID) - if err != nil { - return false - } - return desc.Status == enumspb.WORKFLOW_EXECUTION_STATUS_FAILED + s.EventuallyWithT(func(t *assert.CollectT) { + desc, err := env.SdkClient().DescribeWorkflow(env.Context(), wfID, canRunID) + require.NoError(t, err) + require.Equal(t, enumspb.WORKFLOW_EXECUTION_STATUS_FAILED, desc.Status) }, 10*time.Second, 100*time.Millisecond) // Wait for retry run to complete. var retryRunID string - s.Eventually(func() bool { - desc, err := s.SdkClient().DescribeWorkflowExecution(ctx, wfID, "") - if err != nil { - return false - } + s.EventuallyWithT(func(t *assert.CollectT) { + desc, err := env.SdkClient().DescribeWorkflowExecution(env.Context(), wfID, "") + require.NoError(t, err) retryRunID = desc.GetWorkflowExecutionInfo().GetExecution().GetRunId() - return retryRunID != canRunID && - desc.GetWorkflowExecutionInfo().GetStatus() == enumspb.WORKFLOW_EXECUTION_STATUS_COMPLETED + require.NotEqual(t, canRunID, retryRunID) + require.Equal(t, enumspb.WORKFLOW_EXECUTION_STATUS_COMPLETED, desc.GetWorkflowExecutionInfo().GetStatus()) }, 10*time.Second, 100*time.Millisecond) // Verify: retry run's WFT started should have targetDeploymentVersionChanged=false // because OnStart=v2 (preserved from CaN decision) == target=v2 (still current). - retryHistory := s.GetHistory(s.Namespace().String(), &commonpb.WorkflowExecution{ + retryHistory := env.GetHistory(env.Namespace().String(), &commonpb.WorkflowExecution{ WorkflowId: wfID, RunId: retryRunID, }) @@ -6667,38 +6551,37 @@ func (s *Versioning3Suite) TestRetryOfDeclinedCaN_SignalsOnNewTarget() { // // 5. Set v2 as current again → signal → target=v2, no declined → signal fires (assert true) func (s *Versioning3Suite) TestPinnedCaN_RollbackResetsDeclined() { - ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second) - defer cancel() + env := s.setupEnv() - tv1 := testvars.New(s).WithBuildIDNumber(1) + tv1 := env.Tv().WithBuildIDNumber(1) tv2 := tv1.WithBuildIDNumber(2) // Start async poller for v1 that will handle the first WFT and declare pinned behavior wftCompleted := make(chan struct{}) - s.pollWftAndHandle(tv1, false, wftCompleted, + s.pollWftAndHandle(env, tv1, false, wftCompleted, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) return respondEmptyWft(tv1, false, vbPinned), nil }) // Wait for v1 to be registered, then set it as current - s.waitForDeploymentDataPropagation(tv1, versionStatusInactive, false, tqTypeWf) - s.setCurrentDeployment(tv1) + s.waitForDeploymentDataPropagation(env, tv1, versionStatusInactive, false, tqTypeWf) + s.setCurrentDeployment(env, tv1) // Start workflow — first WFT is handled by the async poller above (pinned on v1) - runID := s.startWorkflow(tv1, nil) + runID := s.startWorkflow(env, tv1, nil) execution := tv1.WithRunID(runID).WorkflowExecution() - s.WaitForChannel(ctx, wftCompleted) - s.verifyWorkflowVersioning(s.Assertions, tv1, vbPinned, tv1.Deployment(), nil, nil) + env.WaitForChannel(env.Context(), wftCompleted) + s.verifyWorkflowVersioning(env, s.Assertions, tv1, vbPinned, tv1.Deployment(), nil, nil) - s.pollUntilRegistered(ctx, tv2) - s.setCurrentDeployment(tv2) + s.pollUntilRegistered(env, tv2) + s.setCurrentDeployment(env, tv2) // Trigger WFT via signal - s.triggerNormalWFT(ctx, tv1, execution) + s.triggerNormalWFT(env, tv1, execution) // Process the WFT: verify targetDeploymentVersionChanged=true, then CaN without AU - s.pollWftAndHandle(tv1, false, nil, + s.pollWftAndHandle(env, tv1, false, nil, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) var lastStarted *historypb.HistoryEvent @@ -6730,7 +6613,7 @@ func (s *Versioning3Suite) TestPinnedCaN_RollbackResetsDeclined() { }) // Step 3: CaN run starts on v1. Target=v2, declined=v2 → case 4 suppresses signal. - s.pollWftAndHandle(tv1, false, nil, + s.pollWftAndHandle(env, tv1, false, nil, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) s.NotEqual(execution.RunId, task.WorkflowExecution.RunId, @@ -6751,13 +6634,13 @@ func (s *Versioning3Suite) TestPinnedCaN_RollbackResetsDeclined() { }) // Step 4: Rollback — set v1 as current again. Target becomes v1. - s.setCurrentDeployment(tv1) - s.waitForDeploymentDataPropagation(tv1, versionStatusCurrent, false, tqTypeWf) + s.setCurrentDeployment(env, tv1) + s.waitForDeploymentDataPropagation(env, tv1, versionStatusCurrent, false, tqTypeWf) // Trigger WFT via signal so case 3 executes (effective=v1 == target=v1) - s.triggerNormalWFT(ctx, tv1, execution) + s.triggerNormalWFT(env, tv1, execution) - s.pollWftAndHandle(tv1, false, nil, + s.pollWftAndHandle(env, tv1, false, nil, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) var lastStarted *historypb.HistoryEvent @@ -6773,15 +6656,15 @@ func (s *Versioning3Suite) TestPinnedCaN_RollbackResetsDeclined() { }) // Step 5: Set v2 as current again. Target goes back to v2. - s.setCurrentDeployment(tv2) - s.waitForDeploymentDataPropagation(tv2, versionStatusCurrent, false, tqTypeWf) + s.setCurrentDeployment(env, tv2) + s.waitForDeploymentDataPropagation(env, tv2, versionStatusCurrent, false, tqTypeWf) // Trigger WFT via signal - s.triggerNormalWFT(ctx, tv1, execution) + s.triggerNormalWFT(env, tv1, execution) // This is the key assertion: declined should have been cleared in step 4, // so target=v2 should fire the signal again. - s.pollWftAndHandle(tv1, false, nil, + s.pollWftAndHandle(env, tv1, false, nil, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) var lastStarted *historypb.HistoryEvent @@ -6813,33 +6696,32 @@ func (s *Versioning3Suite) TestPinnedCaN_RollbackResetsDeclined() { // 3. Unset current deployment → target becomes nil // 4. New run's WFT: targetDeploymentVersionChanged=true (nil target != v1) func (s *Versioning3Suite) TestPinnedCaN_NeverSignaled_NewRunGetsSignalForUnversioned() { - ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second) - defer cancel() + env := s.setupEnv() - tv1 := testvars.New(s).WithBuildIDNumber(1) + tv1 := env.Tv().WithBuildIDNumber(1) // Start async poller for v1 that will handle the first WFT and declare pinned behavior wftCompleted := make(chan struct{}) - s.pollWftAndHandle(tv1, false, wftCompleted, + s.pollWftAndHandle(env, tv1, false, wftCompleted, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) return respondEmptyWft(tv1, false, vbPinned), nil }) // Wait for v1 to be registered, then set it as current - s.waitForDeploymentDataPropagation(tv1, versionStatusInactive, false, tqTypeWf) - s.setCurrentDeployment(tv1) + s.waitForDeploymentDataPropagation(env, tv1, versionStatusInactive, false, tqTypeWf) + s.setCurrentDeployment(env, tv1) // Start workflow — first WFT is handled by the async poller above (pinned on v1) - runID := s.startWorkflow(tv1, nil) + runID := s.startWorkflow(env, tv1, nil) execution := tv1.WithRunID(runID).WorkflowExecution() - s.WaitForChannel(ctx, wftCompleted) - s.verifyWorkflowVersioning(s.Assertions, tv1, vbPinned, tv1.Deployment(), nil, nil) + env.WaitForChannel(env.Context(), wftCompleted) + s.verifyWorkflowVersioning(env, s.Assertions, tv1, vbPinned, tv1.Deployment(), nil, nil) // Trigger a WFT and CaN without any target change — no signal fires. // This means LastNotifiedTargetVersion is never set (nil wrapper). - s.triggerNormalWFT(ctx, tv1, execution) - s.pollWftAndHandle(tv1, false, nil, + s.triggerNormalWFT(env, tv1, execution) + s.pollWftAndHandle(env, tv1, false, nil, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) // Verify no signal — target hasn't changed. @@ -6868,12 +6750,12 @@ func (s *Versioning3Suite) TestPinnedCaN_NeverSignaled_NewRunGetsSignalForUnvers }) // Unset the current deployment — target becomes nil (unversioned). - s.unsetCurrentDeployment(tv1) + s.unsetCurrentDeployment(env, tv1) // New run starts on v1 (pinned, inherited). Since LastNotifiedTargetVersion // was nil (wrapper absent) on the previous run, NotificationSuppressedTargetVersion // on this run is nil. Target=nil != effective=v1 → signal fires. - s.pollWftAndHandle(tv1, false, nil, + s.pollWftAndHandle(env, tv1, false, nil, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) s.NotEqual(execution.RunId, task.WorkflowExecution.RunId, @@ -6903,31 +6785,28 @@ func (s *Versioning3Suite) TestPinnedCaN_NeverSignaled_NewRunGetsSignalForUnvers // 3. Signal → WFT has targetDeploymentVersionChanged=true (v1 != nil) // 4. CaN with AUTO_UPGRADE → new run picked up by unversioned poller → completes func (s *Versioning3Suite) TestPinnedCaN_UpgradeToUnversioned() { - s.RunTestWithMatchingBehavior(func() { - ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second) - defer cancel() - - tv1 := testvars.New(s).WithBuildIDNumber(1) - execution, _ := s.drainWorkflowTaskAfterSetCurrent(tv1) + s.runTestWithMatchingBehavior(func(env *testcore.TestEnv, s *Versioning3Suite) { + tv1 := env.Tv().WithBuildIDNumber(1) + execution, _ := s.drainWorkflowTaskAfterSetCurrent(env, tv1) // Trigger a normal WFT and declare pinned behavior to make the workflow pinned on v1. - s.triggerNormalWFT(ctx, tv1, execution) - s.pollWftAndHandle(tv1, false, nil, + s.triggerNormalWFT(env, tv1, execution) + s.pollWftAndHandle(env, tv1, false, nil, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) return respondEmptyWft(tv1, false, vbPinned), nil }) - s.verifyWorkflowVersioning(s.Assertions, tv1, vbPinned, tv1.Deployment(), nil, nil) + s.verifyWorkflowVersioning(env, s.Assertions, tv1, vbPinned, tv1.Deployment(), nil, nil) // Unset the current deployment — target becomes nil (unversioned). - s.unsetCurrentDeployment(tv1) + s.unsetCurrentDeployment(env, tv1) // Signal to trigger a new WFT. - s.triggerNormalWFT(ctx, tv1, execution) + s.triggerNormalWFT(env, tv1, execution) // Poll as v1 (pinned workflow still dispatched to v1 poller), verify the signal, // and issue CaN with AUTO_UPGRADE initial behavior. - s.pollWftAndHandle(tv1, false, nil, + s.pollWftAndHandle(env, tv1, false, nil, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) @@ -6964,7 +6843,7 @@ func (s *Versioning3Suite) TestPinnedCaN_UpgradeToUnversioned() { // Poll unversioned to pick up the CaN run's first WFT and complete the workflow. wftNewRunDone := make(chan struct{}) - s.unversionedPollWftAndHandle(tv1, false, wftNewRunDone, + s.unversionedPollWftAndHandle(env, tv1, false, wftNewRunDone, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) s.Equal(execution.GetWorkflowId(), task.WorkflowExecution.GetWorkflowId(), @@ -6973,30 +6852,31 @@ func (s *Versioning3Suite) TestPinnedCaN_UpgradeToUnversioned() { "CaN run should have a different run ID") return respondCompleteWorkflowUnversioned(tv1), nil }) - s.WaitForChannel(ctx, wftNewRunDone) + env.WaitForChannel(env.Context(), wftNewRunDone) }) } func (s *Versioning3Suite) TestVersioning3_NoWorkerVersionOnStartedEvents() { - tv := testvars.New(s) + env := s.setupEnv() + tv := env.Tv() tvUpd1 := tv.WithUpdateIDNumber(1).WithMessageIDNumber(1) tvUpd2 := tv.WithUpdateIDNumber(2) stamp := &commonpb.WorkerVersionStamp{UseVersioning: true, BuildId: tv.BuildID()} // Start workflow and drain first WFT. - execution, _ := s.drainWorkflowTaskAfterSetCurrent(tv) + execution, _ := s.drainWorkflowTaskAfterSetCurrent(env, tv) // Send first update. - updateResultCh := sendUpdateNoError(s, tvUpd1) + updateResultCh := sendUpdateNoError(env, tvUpd1) // Poll and handle: accept + complete the first update. - s.pollWftAndHandle(tv, false, nil, + s.pollWftAndHandle(env, tv, false, nil, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) s.NotEmpty(task.Messages) return &workflowservice.RespondWorkflowTaskCompletedRequest{ - Commands: s.UpdateAcceptCompleteCommands(tvUpd1), - Messages: s.UpdateAcceptCompleteMessages(tvUpd1, task.Messages[0]), + Commands: s.updateUtils().UpdateAcceptCompleteCommands(tvUpd1), + Messages: s.updateUtils().UpdateAcceptCompleteMessages(tvUpd1, task.Messages[0]), WorkerVersionStamp: stamp, VersioningBehavior: vbUnpinned, DeploymentOptions: &deploymentpb.WorkerDeploymentOptions{ @@ -7009,20 +6889,20 @@ func (s *Versioning3Suite) TestVersioning3_NoWorkerVersionOnStartedEvents() { <-updateResultCh // Send second update. - updateResultCh2 := sendUpdate(context.Background(), s, tvUpd2) + updateResultCh2 := sendUpdate(env.Context(), env, tvUpd2) // Poll and fail the WFT with deployment options to trigger a transient retry. failCtx := testcore.NewContext() - pollResp, err := s.FrontendClient().PollWorkflowTaskQueue(failCtx, &workflowservice.PollWorkflowTaskQueueRequest{ - Namespace: s.Namespace().String(), + pollResp, err := env.FrontendClient().PollWorkflowTaskQueue(failCtx, &workflowservice.PollWorkflowTaskQueueRequest{ + Namespace: env.Namespace().String(), TaskQueue: tv.TaskQueue(), Identity: tv.WorkerIdentity(), DeploymentOptions: tv.WorkerDeploymentOptions(true), }) s.NoError(err) s.NotEmpty(pollResp.GetTaskToken()) - _, err = s.FrontendClient().RespondWorkflowTaskFailed(failCtx, &workflowservice.RespondWorkflowTaskFailedRequest{ - Namespace: s.Namespace().String(), + _, err = env.FrontendClient().RespondWorkflowTaskFailed(failCtx, &workflowservice.RespondWorkflowTaskFailedRequest{ + Namespace: env.Namespace().String(), TaskToken: pollResp.TaskToken, Cause: enumspb.WORKFLOW_TASK_FAILED_CAUSE_WORKFLOW_WORKER_UNHANDLED_FAILURE, Identity: tv.WorkerIdentity(), @@ -7032,7 +6912,7 @@ func (s *Versioning3Suite) TestVersioning3_NoWorkerVersionOnStartedEvents() { s.NoError(err) // Poll the retried (transient) WFT and fail the workflow. - s.pollWftAndHandle(tv, false, nil, + s.pollWftAndHandle(env, tv, false, nil, func(task *workflowservice.PollWorkflowTaskQueueResponse) (*workflowservice.RespondWorkflowTaskCompletedRequest, error) { s.NotNil(task) return &workflowservice.RespondWorkflowTaskCompletedRequest{ @@ -7059,7 +6939,7 @@ func (s *Versioning3Suite) TestVersioning3_NoWorkerVersionOnStartedEvents() { <-updateResultCh2 // Verify: no WorkflowTaskStarted event should have WorkerVersion set. - events := s.GetHistory(s.Namespace().String(), execution) + events := env.GetHistory(env.Namespace().String(), execution) for _, event := range events { if event.GetEventType() == enumspb.EVENT_TYPE_WORKFLOW_TASK_STARTED { s.Nil( @@ -7073,11 +6953,7 @@ func (s *Versioning3Suite) TestVersioning3_NoWorkerVersionOnStartedEvents() { } func (s *Versioning3Suite) skipBeforeVersion(version workerdeployment.DeploymentWorkflowVersion) { - if s.deploymentWorkflowVersion < version { + if versioning3DeploymentWorkflowVersion < version { s.T().Skipf("test supports workflow version %v and newer", version) } } - -func (s *Versioning3Suite) Context() context.Context { - return s.T().Context() -} diff --git a/tests/versioning_test.go b/tests/versioning_test.go index 762a0d5cad..3993862c3f 100644 --- a/tests/versioning_test.go +++ b/tests/versioning_test.go @@ -13,6 +13,8 @@ import ( "time" "github.com/dgryski/go-farm" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" "github.com/stretchr/testify/suite" commandpb "go.temporal.io/api/command/v1" commonpb "go.temporal.io/api/common/v1" @@ -3953,14 +3955,14 @@ func (s *VersioningIntegSuite) TestDescribeTaskQueueEnhanced_Versioned_Reachabil s.WaitForChannel(ctx, started) // 2. Wait for visibility to show A as running with BuildId SearchAttribute 'assigned:A' - s.Eventually(func() bool { + s.EventuallyWithT(func(t *assert.CollectT) { queryARunning := fmt.Sprintf("TaskQueue = '%s' AND BuildIds IN ('assigned:A') AND ExecutionStatus = \"Running\"", tq) resp, err := s.FrontendClient().CountWorkflowExecutions(ctx, &workflowservice.CountWorkflowExecutionsRequest{ Namespace: s.Namespace().String(), Query: queryARunning, }) - s.NoError(err) - return resp.GetCount() > 0 + require.NoError(t, err) + require.Positive(t, resp.GetCount()) }, 5*time.Second, 50*time.Millisecond) // 3. Commit a different build id --> A should now only be reachable via visibility query @@ -3975,10 +3977,10 @@ func (s *VersioningIntegSuite) TestDescribeTaskQueueEnhanced_Versioned_Reachabil s.NoError(s.SdkClient().SignalWorkflow(ctx, run.GetID(), "", "wait", nil)) // 6. Query reachability(A) --> eventually shows closed_only by visibility db (after TTL passes and A is closed in visibility) - s.Eventually(func() bool { - return s.checkBuildIdReachability(ctx, tq, &taskqueuepb.TaskQueueVersionSelection{BuildIds: []string{"A"}}, map[string]enumspb.BuildIdTaskReachability{ + s.EventuallyWithT(func(t *assert.CollectT) { + require.True(t, s.checkBuildIdReachability(ctx, tq, &taskqueuepb.TaskQueueVersionSelection{BuildIds: []string{"A"}}, map[string]enumspb.BuildIdTaskReachability{ "A": enumspb.BUILD_ID_TASK_REACHABILITY_CLOSED_WORKFLOWS_ONLY, // closed_only by visibility db (after TTL) - }) + })) }, 5*time.Second, 50*time.Millisecond) } @@ -4021,14 +4023,14 @@ func (s *VersioningIntegSuite) TestDescribeTaskQueueEnhanced_Versioned_BasicReac s.WaitForChannel(ctx, started) // wait for visibility to show A as running with BuildId SearchAttribute 'assigned:A' - s.Eventually(func() bool { + s.EventuallyWithT(func(t *assert.CollectT) { queryARunning := fmt.Sprintf("TaskQueue = '%s' AND BuildIds IN ('assigned:A') AND ExecutionStatus = \"Running\"", tq) resp, err := s.FrontendClient().CountWorkflowExecutions(ctx, &workflowservice.CountWorkflowExecutionsRequest{ Namespace: s.Namespace().String(), Query: queryARunning, }) - s.NoError(err) - return resp.GetCount() > 0 + require.NoError(t, err) + require.Positive(t, resp.GetCount()) }, 3*time.Second, 50*time.Millisecond) // commit a different build ID --> A should now only be reachable via visibility query, B reachable as default @@ -4045,11 +4047,11 @@ func (s *VersioningIntegSuite) TestDescribeTaskQueueEnhanced_Versioned_BasicReac s.NoError(s.SdkClient().SignalWorkflow(ctx, run.GetID(), "", "wait", nil)) // Query reachability(A) --> eventually shows closed_only by visibility db (after TTL passes and A is closed in visibility) - s.Eventually(func() bool { - return s.checkBuildIdReachability(ctx, tq, &taskqueuepb.TaskQueueVersionSelection{BuildIds: []string{"A"}}, map[string]enumspb.BuildIdTaskReachability{ + s.EventuallyWithT(func(t *assert.CollectT) { + require.True(t, s.checkBuildIdReachability(ctx, tq, &taskqueuepb.TaskQueueVersionSelection{BuildIds: []string{"A"}}, map[string]enumspb.BuildIdTaskReachability{ "A": enumspb.BUILD_ID_TASK_REACHABILITY_CLOSED_WORKFLOWS_ONLY, // closed_only by visibility db (after TTL) "B": enumspb.BUILD_ID_TASK_REACHABILITY_REACHABLE, // reachable by default assignment rule - }) + })) }, 5*time.Second, 50*time.Millisecond) } @@ -4073,7 +4075,7 @@ func (s *VersioningIntegSuite) TestDescribeTaskQueueEnhanced_Unversioned() { workerMap[wId] = w } - s.Eventually(func() bool { + s.EventuallyWithT(func(t *assert.CollectT) { resp, err := s.FrontendClient().DescribeTaskQueue(ctx, &workflowservice.DescribeTaskQueueRequest{ Namespace: s.Namespace().String(), TaskQueue: &taskqueuepb.TaskQueue{Name: tq, Kind: enumspb.TASK_QUEUE_KIND_NORMAL}, @@ -4084,11 +4086,11 @@ func (s *VersioningIntegSuite) TestDescribeTaskQueueEnhanced_Unversioned() { ReportTaskReachability: true, ReportStats: false, }) - s.NoError(err) - s.NotNil(resp) - s.Len(resp.GetVersionsInfo(), 1, "should be 1 because only default/unversioned queue") //nolint:staticcheck + require.NoError(t, err) + require.NotNil(t, resp) + require.Len(t, resp.GetVersionsInfo(), 1, "should be 1 because only default/unversioned queue") //nolint:staticcheck versionInfo := resp.GetVersionsInfo()[""] - s.Equal(enumspb.BUILD_ID_TASK_REACHABILITY_REACHABLE, versionInfo.GetTaskReachability()) + require.Equal(t, enumspb.BUILD_ID_TASK_REACHABILITY_REACHABLE, versionInfo.GetTaskReachability()) var pollersInfo []*taskqueuepb.PollerInfo for _, t := range versionInfo.GetTypesInfo() { pollersInfo = append(pollersInfo, t.GetPollers()...) @@ -4096,7 +4098,7 @@ func (s *VersioningIntegSuite) TestDescribeTaskQueueEnhanced_Unversioned() { foundN := 0 for wId := range workerMap { for _, pi := range pollersInfo { - s.False(pi.GetWorkerVersionCapabilities().GetUseVersioning()) + require.False(t, pi.GetWorkerVersionCapabilities().GetUseVersioning()) if pi.GetIdentity() == wId { foundN++ break @@ -4104,7 +4106,7 @@ func (s *VersioningIntegSuite) TestDescribeTaskQueueEnhanced_Unversioned() { } } - return foundN == workerN + require.Equal(t, workerN, foundN) }, 3*time.Second, 50*time.Millisecond) } @@ -4124,7 +4126,7 @@ func (s *VersioningIntegSuite) TestDescribeTaskQueueEnhanced_ReportFlags() { defer w.Stop() // wait for pollers to show up, verify both ReportPollers and ReportTaskReachability - s.Eventually(func() bool { + s.EventuallyWithT(func(t *assert.CollectT) { resp, err := s.FrontendClient().DescribeTaskQueue(ctx, &workflowservice.DescribeTaskQueueRequest{ Namespace: s.Namespace().String(), TaskQueue: &taskqueuepb.TaskQueue{Name: tq, Kind: enumspb.TASK_QUEUE_KIND_NORMAL}, @@ -4134,23 +4136,25 @@ func (s *VersioningIntegSuite) TestDescribeTaskQueueEnhanced_ReportFlags() { ReportPollers: true, ReportTaskReachability: true, }) - s.NoError(err) - s.NotNil(resp) - s.Len(resp.GetVersionsInfo(), 1, "should be 1 because only default/unversioned queue") //nolint:staticcheck + require.NoError(t, err) + require.NotNil(t, resp) + require.Len(t, resp.GetVersionsInfo(), 1, "should be 1 because only default/unversioned queue") //nolint:staticcheck versionInfo := resp.GetVersionsInfo()[""] - s.Equal(enumspb.BUILD_ID_TASK_REACHABILITY_REACHABLE, versionInfo.GetTaskReachability()) + require.Equal(t, enumspb.BUILD_ID_TASK_REACHABILITY_REACHABLE, versionInfo.GetTaskReachability()) var pollersInfo []*taskqueuepb.PollerInfo for _, t := range versionInfo.GetTypesInfo() { pollersInfo = append(pollersInfo, t.GetPollers()...) } + foundPoller := false for _, pi := range pollersInfo { - s.False(pi.GetWorkerVersionCapabilities().GetUseVersioning()) + require.False(t, pi.GetWorkerVersionCapabilities().GetUseVersioning()) if pi.GetIdentity() == wId { - return true + foundPoller = true + break } } - return false + require.True(t, foundPoller) }, 3*time.Second, 50*time.Millisecond) // ask for reachability only @@ -4269,13 +4273,13 @@ func (s *VersioningIntegSuite) TestDescribeTaskQueueLegacy_VersionSets() { s.NoError(w2.Start()) defer w2.Stop() - s.Eventually(func() bool { + s.EventuallyWithT(func(t *assert.CollectT) { resp, err := s.FrontendClient().DescribeTaskQueue(ctx, &workflowservice.DescribeTaskQueueRequest{ Namespace: s.Namespace().String(), TaskQueue: &taskqueuepb.TaskQueue{Name: tq, Kind: enumspb.TASK_QUEUE_KIND_NORMAL}, TaskQueueType: enumspb.TASK_QUEUE_TYPE_WORKFLOW, }) - s.NoError(err) + require.NoError(t, err) havePoller := func(v string) bool { for _, p := range resp.Pollers { if p.WorkerVersionCapabilities.UseVersioning && v == p.WorkerVersionCapabilities.BuildId { @@ -4285,7 +4289,9 @@ func (s *VersioningIntegSuite) TestDescribeTaskQueueLegacy_VersionSets() { return false } // v1 polls get rejected because v11 is newer - return !havePoller(v1) && havePoller(v11) && havePoller(v2) + require.False(t, havePoller(v1)) + require.True(t, havePoller(v11)) + require.True(t, havePoller(v2)) }, 3*time.Second, 50*time.Millisecond) } @@ -4328,10 +4334,10 @@ func (s *VersioningIntegSuite) TestDescribeWorkflowExecution() { s.WaitForChannel(ctx, started1) // describe and check build ID - s.Eventually(func() bool { + s.EventuallyWithT(func(t *assert.CollectT) { resp, err := s.SdkClient().DescribeWorkflowExecution(ctx, run.GetID(), "") - s.NoError(err) - return v1 == resp.GetWorkflowExecutionInfo().GetMostRecentWorkerVersionStamp().GetBuildId() + require.NoError(t, err) + require.Equal(t, v1, resp.GetWorkflowExecutionInfo().GetMostRecentWorkerVersionStamp().GetBuildId()) }, 5*time.Second, 100*time.Millisecond) // now register v11 as newer compatible with v1 @@ -4356,10 +4362,10 @@ func (s *VersioningIntegSuite) TestDescribeWorkflowExecution() { s.NoError(s.SdkClient().SignalWorkflow(ctx, run.GetID(), "", "wait", nil)) s.WaitForChannel(ctx, started11) - s.Eventually(func() bool { + s.EventuallyWithT(func(t *assert.CollectT) { resp, err := s.SdkClient().DescribeWorkflowExecution(ctx, run.GetID(), "") - s.NoError(err) - return v11 == resp.GetWorkflowExecutionInfo().GetMostRecentWorkerVersionStamp().GetBuildId() + require.NoError(t, err) + require.Equal(t, v11, resp.GetWorkflowExecutionInfo().GetMostRecentWorkerVersionStamp().GetBuildId()) }, 5*time.Second, 100*time.Millisecond) // unblock. it should complete @@ -4901,10 +4907,10 @@ func (s *VersioningIntegSuite) waitForPropagation( remaining[partAndType{i, enumspb.TASK_QUEUE_TYPE_ACTIVITY}] = struct{}{} remaining[partAndType{i, enumspb.TASK_QUEUE_TYPE_WORKFLOW}] = struct{}{} } - s.Eventually(func() bool { + s.EventuallyWithT(func(t *assert.CollectT) { for pt := range remaining { f, err := tqid.NewTaskQueueFamily(s.NamespaceID().String(), taskQueue) - s.NoError(err) + require.NoError(t, err) partition := f.TaskQueue(pt.tp).NormalPartition(pt.part) // Use lower-level GetTaskQueueUserData instead of GetWorkerBuildIdCompatibility // here so that we can target activity queues. @@ -4915,12 +4921,12 @@ func (s *VersioningIntegSuite) waitForPropagation( TaskQueue: partition.RpcName(), TaskQueueType: partition.TaskType(), }) - s.NoError(err) + require.NoError(t, err) if condition(res.GetUserData().GetData().GetVersioningData()) { delete(remaining, pt) } } - return len(remaining) == 0 + require.Empty(t, remaining) }, 10*time.Second, 100*time.Millisecond) }