@@ -87,7 +87,6 @@ func TestReplicationSimulation(t *testing.T) {
8787 startTime := time .Now ().UTC ()
8888 simTypes .Logf (t , "Simulation start time: %v" , startTime )
8989 for i , op := range simCfg .Operations {
90- op := op
9190 waitForOpTime (t , op , startTime )
9291 var err error
9392 switch op .Type {
@@ -100,9 +99,9 @@ func TestReplicationSimulation(t *testing.T) {
10099 case simTypes .ReplicationSimulationOperationValidate :
101100 err = validate (t , op , simCfg , sim )
102101 case simTypes .ReplicationSimulationOperationQueryWorkflow :
103- err = queryWorkflow (t , op , simCfg )
102+ err = queryWorkflow (t , op , simCfg , sim )
104103 case simTypes .ReplicationSimulationOperationSignalWithStartWorkflow :
105- err = signalWithStartWorkflow (t , op , simCfg )
104+ err = signalWithStartWorkflow (t , op , simCfg , sim )
106105 case simTypes .ReplicationSimulationOperationMigrateDomainToActiveActive :
107106 err = migrateDomainToActiveActive (t , op , simCfg )
108107 case simTypes .ReplicationSimulationOperationValidateWorkflowReplication :
@@ -330,7 +329,7 @@ func migrateDomainToActiveActive(t *testing.T, op *simTypes.Operation, simCfg *s
330329 return nil
331330}
332331
333- func queryWorkflow (t * testing.T , op * simTypes.Operation , simCfg * simTypes.ReplicationSimulationConfig ) error {
332+ func queryWorkflow (t * testing.T , op * simTypes.Operation , simCfg * simTypes.ReplicationSimulationConfig , sim * simTypes. ReplicationSimulation ) error {
334333 t .Helper ()
335334
336335 simTypes .Logf (t , "Querying workflow: %s on domain %s on cluster: %s" , op .WorkflowID , op .Domain , op .Cluster )
@@ -344,11 +343,22 @@ func queryWorkflow(t *testing.T, op *simTypes.Operation, simCfg *simTypes.Replic
344343 consistencyLevel = types .QueryConsistencyLevelStrong .Ptr ()
345344 }
346345
346+ // Prepare workflow execution - use specific RunID if provided via runIDKey
347+ executionRequest := & types.WorkflowExecution {
348+ WorkflowID : op .WorkflowID ,
349+ }
350+ if op .RunIDKey != "" {
351+ if runID , err := sim .GetRunID (op .RunIDKey ); err == nil && runID != "" {
352+ executionRequest .RunID = runID
353+ simTypes .Logf (t , "Using stored RunID %s for query (key: %s)" , runID , op .RunIDKey )
354+ } else {
355+ return fmt .Errorf ("runIDKey %s specified but no RunID found in registry" , op .RunIDKey )
356+ }
357+ }
358+
347359 queryResp , err := frontendCl .QueryWorkflow (ctx , & types.QueryWorkflowRequest {
348- Domain : op .Domain ,
349- Execution : & types.WorkflowExecution {
350- WorkflowID : op .WorkflowID ,
351- },
360+ Domain : op .Domain ,
361+ Execution : executionRequest ,
352362 Query : & types.WorkflowQuery {
353363 QueryType : op .Query ,
354364 },
@@ -373,6 +383,7 @@ func signalWithStartWorkflow(
373383 t * testing.T ,
374384 op * simTypes.Operation ,
375385 simCfg * simTypes.ReplicationSimulationConfig ,
386+ sim * simTypes.ReplicationSimulation ,
376387) error {
377388 t .Helper ()
378389 simTypes .Logf (t , "SignalWithStart workflow: %s on domain %s on cluster: %s" , op .WorkflowID , op .Domain , op .Cluster )
@@ -398,7 +409,15 @@ func signalWithStartWorkflow(
398409 return err
399410 }
400411
401- simTypes .Logf (t , "SignalWithStart workflow: %s on domain %s on cluster: %s. RunID: %s" , op .WorkflowID , op .Domain , op .Cluster , signalResp .GetRunID ())
412+ runID := signalResp .GetRunID ()
413+ simTypes .Logf (t , "SignalWithStart workflow: %s on domain %s on cluster: %s. RunID: %s" , op .WorkflowID , op .Domain , op .Cluster , runID )
414+
415+ // Store RunID if runIDKey is specified
416+ if op .RunIDKey != "" {
417+ sim .StoreRunID (op .RunIDKey , runID )
418+ simTypes .Logf (t , "Stored RunID %s with key: %s" , runID , op .RunIDKey )
419+ }
420+
402421 return nil
403422}
404423
@@ -415,73 +434,89 @@ func validate(
415434
416435 simTypes .Logf (t , "Validating workflow: %s on cluster: %s" , op .WorkflowID , op .Cluster )
417436
437+ consistencyLevel := types .QueryConsistencyLevelEventual .Ptr ()
438+ if op .ConsistencyLevel == "strong" {
439+ consistencyLevel = types .QueryConsistencyLevelStrong .Ptr ()
440+ }
441+
418442 ctx , cancel := context .WithTimeout (context .Background (), 2 * time .Second )
419443 defer cancel ()
420444
421445 // Prepare workflow execution - use specific RunID if provided via runIDKey
422- execution := & types.WorkflowExecution {
446+ executionRequest := & types.WorkflowExecution {
423447 WorkflowID : op .WorkflowID ,
424448 }
425449 if op .RunIDKey != "" {
426450 if runID , err := sim .GetRunID (op .RunIDKey ); err == nil && runID != "" {
427- execution .RunID = runID
451+ executionRequest .RunID = runID
428452 simTypes .Logf (t , "Using stored RunID %s for validation (key: %s)" , runID , op .RunIDKey )
429453 } else {
430454 return fmt .Errorf ("runIDKey %s specified but no RunID found in registry" , op .RunIDKey )
431455 }
432456 }
457+
433458 resp , err := simCfg .MustGetFrontendClient (t , op .Cluster ).DescribeWorkflowExecution (ctx ,
434459 & types.DescribeWorkflowExecutionRequest {
435- Domain : op .Domain ,
436- Execution : execution ,
460+ Domain : op .Domain ,
461+ Execution : executionRequest ,
462+ QueryConsistencyLevel : consistencyLevel ,
437463 })
438464 if err != nil {
439465 return err
440466 }
441467
468+ workflowStatus := resp .GetWorkflowExecutionInfo ().GetCloseStatus ()
469+ workflowCloseTime := resp .GetWorkflowExecutionInfo ().GetCloseTime ()
442470 switch op .Want .Status {
443471 case "completed" :
444472 // Validate workflow completed
445- if resp . GetWorkflowExecutionInfo (). GetCloseStatus () != types .WorkflowExecutionCloseStatusCompleted || resp . GetWorkflowExecutionInfo (). GetCloseTime () == 0 {
446- return fmt .Errorf ("workflow %s not completed. status: %s, close time: %v" , op .WorkflowID , resp . GetWorkflowExecutionInfo (). GetCloseStatus () , time .Unix (0 , resp . GetWorkflowExecutionInfo (). GetCloseTime () ))
473+ if workflowStatus != types .WorkflowExecutionCloseStatusCompleted || workflowCloseTime == 0 {
474+ return fmt .Errorf ("workflow %s not completed. status: %s, close time: %v" , op .WorkflowID , workflowStatus , time .Unix (0 , workflowCloseTime ))
447475 }
448476 case "failed" :
449477 // Validate workflow failed
450- if resp . GetWorkflowExecutionInfo (). GetCloseStatus () != types .WorkflowExecutionCloseStatusFailed || resp . GetWorkflowExecutionInfo (). GetCloseTime () == 0 {
451- return fmt .Errorf ("workflow %s not failed. status: %s, close time: %v" , op .WorkflowID , resp . GetWorkflowExecutionInfo (). GetCloseStatus () , time .Unix (0 , resp . GetWorkflowExecutionInfo (). GetCloseTime () ))
478+ if workflowStatus != types .WorkflowExecutionCloseStatusFailed || workflowCloseTime == 0 {
479+ return fmt .Errorf ("workflow %s not failed. status: %s, close time: %v" , op .WorkflowID , workflowStatus , time .Unix (0 , workflowCloseTime ))
452480 }
453481 case "canceled" :
454482 // Validate workflow canceled
455- if resp . GetWorkflowExecutionInfo (). GetCloseStatus () != types .WorkflowExecutionCloseStatusCanceled || resp . GetWorkflowExecutionInfo (). GetCloseTime () == 0 {
456- return fmt .Errorf ("workflow %s not canceled. status: %s, close time: %v" , op .WorkflowID , resp . GetWorkflowExecutionInfo (). GetCloseStatus () , time .Unix (0 , resp . GetWorkflowExecutionInfo (). GetCloseTime () ))
483+ if workflowStatus != types .WorkflowExecutionCloseStatusCanceled || workflowCloseTime == 0 {
484+ return fmt .Errorf ("workflow %s not canceled. status: %s, close time: %v" , op .WorkflowID , workflowStatus , time .Unix (0 , workflowCloseTime ))
457485 }
458486 case "terminated" :
459487 // Validate workflow terminated
460- if resp . GetWorkflowExecutionInfo (). GetCloseStatus () != types .WorkflowExecutionCloseStatusTerminated || resp . GetWorkflowExecutionInfo (). GetCloseTime () == 0 {
461- return fmt .Errorf ("workflow %s not terminated. status: %s, close time: %v" , op .WorkflowID , resp . GetWorkflowExecutionInfo (). GetCloseStatus () , time .Unix (0 , resp . GetWorkflowExecutionInfo (). GetCloseTime () ))
488+ if workflowStatus != types .WorkflowExecutionCloseStatusTerminated || workflowCloseTime == 0 {
489+ return fmt .Errorf ("workflow %s not terminated. status: %s, close time: %v" , op .WorkflowID , workflowStatus , time .Unix (0 , workflowCloseTime ))
462490 }
463491 case "continued-as-new" :
464492 // Validate workflow continued as new
465- if resp . GetWorkflowExecutionInfo (). GetCloseStatus () != types .WorkflowExecutionCloseStatusContinuedAsNew || resp . GetWorkflowExecutionInfo (). GetCloseTime () == 0 {
466- return fmt .Errorf ("workflow %s not continued as new. status: %s, close time: %v" , op .WorkflowID , resp . GetWorkflowExecutionInfo (). GetCloseStatus () , time .Unix (0 , resp . GetWorkflowExecutionInfo (). GetCloseTime () ))
493+ if workflowStatus != types .WorkflowExecutionCloseStatusContinuedAsNew || workflowCloseTime == 0 {
494+ return fmt .Errorf ("workflow %s not continued as new. status: %s, close time: %v" , op .WorkflowID , workflowStatus , time .Unix (0 , workflowCloseTime ))
467495 }
468496 case "timed-out" :
469497 // Validate workflow timed out
470- if resp . GetWorkflowExecutionInfo (). GetCloseStatus () != types .WorkflowExecutionCloseStatusTimedOut || resp . GetWorkflowExecutionInfo (). GetCloseTime () == 0 {
471- return fmt .Errorf ("workflow %s not timed out. status: %s, close time: %v" , op .WorkflowID , resp . GetWorkflowExecutionInfo (). GetCloseStatus () , time .Unix (0 , resp . GetWorkflowExecutionInfo (). GetCloseTime () ))
498+ if workflowStatus != types .WorkflowExecutionCloseStatusTimedOut || workflowCloseTime == 0 {
499+ return fmt .Errorf ("workflow %s not timed out. status: %s, close time: %v" , op .WorkflowID , workflowStatus , time .Unix (0 , workflowCloseTime ))
472500 }
473501 default :
474502 // Validate workflow is running
475- if resp . GetWorkflowExecutionInfo (). GetCloseTime () != 0 {
476- return fmt .Errorf ("workflow %s not running. status: %s, close time: %v" , op .WorkflowID , resp . GetWorkflowExecutionInfo (). GetCloseStatus () , time .Unix (0 , resp . GetWorkflowExecutionInfo (). GetCloseTime () ))
503+ if workflowCloseTime != 0 {
504+ return fmt .Errorf ("workflow %s not running. status: %s, close time: %v" , op .WorkflowID , workflowStatus , time .Unix (0 , workflowCloseTime ))
477505 }
478506 }
479507
480508 simTypes .Logf (t , "Validated workflow: %s on cluster: %s. Status: %s, CloseTime: %v" , op .WorkflowID , op .Cluster , resp .GetWorkflowExecutionInfo ().GetCloseStatus (), time .Unix (0 , resp .GetWorkflowExecutionInfo ().GetCloseTime ()))
481509
482510 // Get history to validate the worker identity that started and completed the workflow
483511 // Some workflows start in cluster0 and complete in cluster1. This is to validate that
484- history , err := getAllHistory (t , simCfg , op .Cluster , op .Domain , op .WorkflowID )
512+ var runID string
513+ if op .RunIDKey != "" {
514+ runID , err = sim .GetRunID (op .RunIDKey )
515+ if err != nil {
516+ return err
517+ }
518+ }
519+ history , err := getAllHistory (t , simCfg , op .Cluster , op .Domain , op .WorkflowID , runID )
485520 if err != nil {
486521 return err
487522 }
@@ -588,17 +623,23 @@ func waitForOpTime(t *testing.T, op *simTypes.Operation, startTime time.Time) {
588623 simTypes .Logf (t , "Operation time (t + %ds) reached: %v" , int (op .At .Seconds ()), startTime .Add (op .At ))
589624}
590625
591- func getAllHistory (t * testing.T , simCfg * simTypes.ReplicationSimulationConfig , clusterName , domainName , wfID string ) ([]types.HistoryEvent , error ) {
626+ func getAllHistory (t * testing.T , simCfg * simTypes.ReplicationSimulationConfig , clusterName , domainName , wfID , runID string ) ([]types.HistoryEvent , error ) {
592627 frontendCl := simCfg .MustGetFrontendClient (t , clusterName )
593628 var nextPageToken []byte
594629 var history []types.HistoryEvent
630+
631+ executionRequest := & types.WorkflowExecution {
632+ WorkflowID : wfID ,
633+ }
634+ if runID != "" {
635+ executionRequest .RunID = runID
636+ }
637+
595638 for {
596639 ctx , cancel := context .WithTimeout (context .Background (), 2 * time .Second )
597640 response , err := frontendCl .GetWorkflowExecutionHistory (ctx , & types.GetWorkflowExecutionHistoryRequest {
598- Domain : domainName ,
599- Execution : & types.WorkflowExecution {
600- WorkflowID : wfID ,
601- },
641+ Domain : domainName ,
642+ Execution : executionRequest ,
602643 MaximumPageSize : 1000 ,
603644 NextPageToken : nextPageToken ,
604645 WaitForNewEvent : false ,
0 commit comments