@@ -326,6 +326,79 @@ public async Task ScheduleNewOrchestrationInstance_IdProvided_TagsProvided()
326326 await this . RunScheduleNewOrchestrationInstanceAsync ( "test" , "input" , options ) ;
327327 }
328328
329+ [ Theory ]
330+ [ InlineData ( false ) ]
331+ [ InlineData ( true ) ]
332+ public async Task RestartAsync_EndToEnd ( bool restartWithNewInstanceId )
333+ {
334+ string originalInstanceId = "test-instance-id" ;
335+ string orchestratorName = "TestOrchestrator" ;
336+ object input = "test-input" ;
337+ string serializedInput = "\" test-input\" " ;
338+
339+ // Create a completed orchestration state
340+ Core . OrchestrationState originalState = CreateState ( input , "test-output" ) ;
341+ originalState . OrchestrationInstance . InstanceId = originalInstanceId ;
342+ originalState . Name = orchestratorName ;
343+ originalState . OrchestrationStatus = Core . OrchestrationStatus . Completed ;
344+
345+ // Setup the mock to return the original orchestration state
346+ this . orchestrationClient
347+ . Setup ( x => x . GetOrchestrationStateAsync ( originalInstanceId , false ) )
348+ . ReturnsAsync ( new List < Core . OrchestrationState > { originalState } ) ;
349+
350+ // Capture the TaskMessage for verification becasue we will create this message at RestartAsync.
351+ TaskMessage ? capturedMessage = null ;
352+ this . orchestrationClient
353+ . Setup ( x => x . CreateTaskOrchestrationAsync ( It . IsAny < TaskMessage > ( ) ) )
354+ . Callback < TaskMessage > ( msg => capturedMessage = msg )
355+ . Returns ( Task . CompletedTask ) ;
356+
357+ string restartedInstanceId = await this . client . RestartAsync ( originalInstanceId , restartWithNewInstanceId ) ;
358+
359+ if ( restartWithNewInstanceId )
360+ {
361+ restartedInstanceId . Should ( ) . NotBe ( originalInstanceId ) ;
362+ }
363+ else
364+ {
365+ restartedInstanceId . Should ( ) . Be ( originalInstanceId ) ;
366+ }
367+
368+ // Verify that CreateTaskOrchestrationAsync was called
369+ this . orchestrationClient . Verify (
370+ x => x . CreateTaskOrchestrationAsync ( It . IsAny < TaskMessage > ( ) ) ,
371+ Times . Once ) ;
372+
373+ // Verify the captured message details
374+ capturedMessage . Should ( ) . NotBeNull ( ) ;
375+ capturedMessage ! . Event . Should ( ) . BeOfType < ExecutionStartedEvent > ( ) ;
376+
377+ var startedEvent = ( ExecutionStartedEvent ) capturedMessage . Event ;
378+ startedEvent . Name . Should ( ) . Be ( orchestratorName ) ;
379+ startedEvent . Input . Should ( ) . Be ( serializedInput ) ;
380+ // TODO: once we support version at ShimDurableTaskClient, we should check version here.
381+ startedEvent . OrchestrationInstance . InstanceId . Should ( ) . Be ( restartedInstanceId ) ;
382+ startedEvent . OrchestrationInstance . ExecutionId . Should ( ) . NotBeNullOrEmpty ( ) ;
383+ startedEvent . OrchestrationInstance . ExecutionId . Should ( ) . NotBe ( originalState . OrchestrationInstance . ExecutionId ) ;
384+ }
385+
386+ [ Fact ]
387+ public async Task RestartAsync_InstanceNotFound_ThrowsArgumentException ( )
388+ {
389+ string nonExistentInstanceId = "non-existent-instance-id" ;
390+
391+ // Setup the mock to client return empty orchestration state (instance not found)
392+ this . orchestrationClient
393+ . Setup ( x => x . GetOrchestrationStateAsync ( nonExistentInstanceId , false ) )
394+ . ReturnsAsync ( new List < Core . OrchestrationState > ( ) ) ;
395+
396+ // RestartAsync should throw an ArgumentException since the instance is not found
397+ Func < Task > restartAction = ( ) => this . client . RestartAsync ( nonExistentInstanceId ) ;
398+
399+ await restartAction . Should ( ) . ThrowAsync < ArgumentException > ( )
400+ . WithMessage ( $ "*An orchestration with the instanceId { nonExistentInstanceId } was not found*") ;
401+ }
329402
330403 static Core . OrchestrationState CreateState (
331404 object input , object ? output = null , DateTimeOffset start = default )
0 commit comments