@@ -432,46 +432,39 @@ public void testWorkflowsCanUseVersioningOverride() {
432432 == PINNED_OVERRIDE_BEHAVIOR_PINNED ));
433433 }
434434
435- public static class TestWorkerVersioningCanV1 implements TestWorkflows .QueryableWorkflow {
435+ @ WorkflowInterface
436+ public interface ContinueAsNewVersionUpgradeWorkflow {
437+ @ WorkflowMethod
438+ String execute (int attempt );
439+ }
440+
441+ public static class TestWorkerVersioningCanV1 implements ContinueAsNewVersionUpgradeWorkflow {
436442 @ Override
437443 @ WorkflowVersioningBehavior (VersioningBehavior .PINNED )
438- public String execute () {
444+ public String execute (int attempt ) {
445+ if (attempt > 0 ) {
446+ return "v1.0" ;
447+ }
439448 while (!Workflow .getInfo ().isTargetWorkerDeploymentVersionChanged ()) {
440449 Workflow .sleep (java .time .Duration .ofMillis (10 ));
441450 }
442451 ContinueAsNewOptions options =
443452 ContinueAsNewOptions .newBuilder ()
444453 .setInitialVersioningBehavior (InitialVersioningBehavior .AUTO_UPGRADE )
445454 .build ();
446- TestWorkflows . QueryableWorkflow next =
447- Workflow .newContinueAsNewStub (TestWorkflows . QueryableWorkflow .class , options );
448- next .execute ();
455+ ContinueAsNewVersionUpgradeWorkflow next =
456+ Workflow .newContinueAsNewStub (ContinueAsNewVersionUpgradeWorkflow .class , options );
457+ next .execute (attempt + 1 );
449458 throw new RuntimeException ("unreachable" );
450459 }
451-
452- @ Override
453- public void mySignal (String arg ) {}
454-
455- @ Override
456- public String getState () {
457- return "v1-can" ;
458- }
459460 }
460461
461- public static class TestWorkerVersioningCanV2 implements TestWorkflows . QueryableWorkflow {
462+ public static class TestWorkerVersioningCanV2 implements ContinueAsNewVersionUpgradeWorkflow {
462463 @ Override
463464 @ WorkflowVersioningBehavior (VersioningBehavior .PINNED )
464- public String execute () {
465+ public String execute (int attempt ) {
465466 return "v2.0" ;
466467 }
467-
468- @ Override
469- public void mySignal (String arg ) {}
470-
471- @ Override
472- public String getState () {
473- return "v2-can" ;
474- }
475468 }
476469
477470 @ Test
@@ -497,13 +490,12 @@ public void testContinueAsNewWithVersionUpgrade() {
497490 waitForRoutingConfigPropagation (v1 );
498491
499492 // Start workflow on v1
500- TestWorkflows . QueryableWorkflow wf =
493+ ContinueAsNewVersionUpgradeWorkflow wf =
501494 testWorkflowRule .newWorkflowStubTimeoutOptions (
502- TestWorkflows . QueryableWorkflow .class , "can-version-upgrade" );
503- WorkflowExecution we = WorkflowClient .start (wf ::execute );
495+ ContinueAsNewVersionUpgradeWorkflow .class , "can-version-upgrade" );
496+ WorkflowExecution we = WorkflowClient .start (wf ::execute , 0 );
504497
505498 // Verify workflow is running on v1
506- Assert .assertEquals ("v1-can" , wf .getState ());
507499 waitForWorkflowRunningOnVersion (we .getWorkflowId (), "1.0" );
508500
509501 // Set v2 as current — triggers targetWorkerDeploymentVersionChanged
0 commit comments