11using  ModelContextProtocol . Client ; 
2- using  ModelContextProtocol . Configuration ; 
3- using  ModelContextProtocol . Protocol . Messages ; 
4- using  ModelContextProtocol . Protocol . Transport ; 
5- using  ModelContextProtocol . Protocol . Types ; 
62using  Microsoft . Extensions . AI ; 
73using  OpenAI ; 
4+ using  ModelContextProtocol . Protocol . Types ; 
5+ using  ModelContextProtocol . Protocol . Messages ; 
86using  System . Text . Json ; 
7+ using  ModelContextProtocol . Configuration ; 
8+ using  ModelContextProtocol . Protocol . Transport ; 
9+ using  Xunit . Sdk ; 
910
1011namespace  ModelContextProtocol . Tests ; 
1112
@@ -61,8 +62,8 @@ public async Task ListTools_Stdio(string clientId)
6162
6263        // act 
6364        await  using  var  client  =  await  _fixture . CreateClientAsync ( clientId ) ; 
64-         var  tools  =  await  client . ListToolsAsync ( ) . ToListAsync ( ) ; 
65-         var  aiFunctions  =  await  client . GetAIFunctionsAsync ( ) ; 
65+         var  tools  =  await  client . ListToolsAsync ( TestContext . Current . CancellationToken ) . ToListAsync ( TestContext . Current . CancellationToken ) ; 
66+         var  aiFunctions  =  await  client . GetAIFunctionsAsync ( TestContext . Current . CancellationToken ) ; 
6667
6768        // assert 
6869        Assert . NotEmpty ( tools ) ; 
@@ -102,9 +103,9 @@ public async Task CallTool_Stdio_ViaAIFunction_EchoServer(string clientId)
102103
103104        // act 
104105        await  using  var  client  =  await  _fixture . CreateClientAsync ( clientId ) ; 
105-         var  aiFunctions  =  await  client . GetAIFunctionsAsync ( ) ; 
106+         var  aiFunctions  =  await  client . GetAIFunctionsAsync ( TestContext . Current . CancellationToken ) ; 
106107        var  echo  =  aiFunctions . Single ( t =>  t . Name  ==  "echo" ) ; 
107-         var  result  =  await  echo . InvokeAsync ( [ new  KeyValuePair < string ,  object ? > ( "message" ,  "Hello MCP!" ) ] ) ; 
108+         var  result  =  await  echo . InvokeAsync ( [ new  KeyValuePair < string ,  object ? > ( "message" ,  "Hello MCP!" ) ] ,   TestContext . Current . CancellationToken ) ; 
108109
109110        // assert 
110111        Assert . NotNull ( result ) ; 
@@ -119,7 +120,7 @@ public async Task ListPrompts_Stdio(string clientId)
119120
120121        // act 
121122        await  using  var  client  =  await  _fixture . CreateClientAsync ( clientId ) ; 
122-         var  prompts  =  await  client . ListPromptsAsync ( ) . ToListAsync ( ) ; 
123+         var  prompts  =  await  client . ListPromptsAsync ( TestContext . Current . CancellationToken ) . ToListAsync ( TestContext . Current . CancellationToken ) ; 
123124
124125        // assert 
125126        Assert . NotEmpty ( prompts ) ; 
@@ -251,7 +252,7 @@ public async Task SubscribeResource_Stdio()
251252        await  client . SubscribeToResourceAsync ( "test://static/resource/1" ,  CancellationToken . None ) ; 
252253
253254        // notifications happen every 5 seconds, so we wait for 10 seconds to ensure we get at least one notification 
254-         await  Task . Delay ( 10000 ) ; 
255+         await  Task . Delay ( 10000 ,   TestContext . Current . CancellationToken ) ; 
255256
256257        // assert 
257258        Assert . True ( counter  >  0 ) ; 
@@ -276,17 +277,17 @@ public async Task UnsubscribeResource_Stdio()
276277        await  client . SubscribeToResourceAsync ( "test://static/resource/1" ,  CancellationToken . None ) ; 
277278
278279        // notifications happen every 5 seconds, so we wait for 10 seconds to ensure we get at least one notification 
279-         await  Task . Delay ( 10000 ) ; 
280+         await  Task . Delay ( 10000 ,   TestContext . Current . CancellationToken ) ; 
280281
281282        // reset counter 
282283        int  counterAfterSubscribe  =  counter ; 
283-          
284+ 
284285        // unsubscribe 
285286        await  client . UnsubscribeFromResourceAsync ( "test://static/resource/1" ,  CancellationToken . None ) ; 
286287        counter  =  0 ; 
287288
288289        // notifications happen every 5 seconds, so we wait for 10 seconds to ensure we would've gotten at least one notification 
289-         await  Task . Delay ( 10000 ) ; 
290+         await  Task . Delay ( 10000 ,   TestContext . Current . CancellationToken ) ; 
290291
291292        // assert 
292293        Assert . True ( counterAfterSubscribe  >  0 ) ; 
@@ -340,7 +341,7 @@ public async Task GetCompletion_Stdio_PromptReference(string clientId)
340341    [ Theory ] 
341342    [ MemberData ( nameof ( GetClients ) ) ] 
342343    public  async  Task  Sampling_Stdio ( string  clientId ) 
343-     {     
344+     { 
344345        // Set up the sampling handler 
345346        int  samplingHandlerCalls  =  0 ; 
346347        await  using  var  client  =  await  _fixture . CreateClientAsync ( clientId ,  new ( ) 
@@ -375,8 +376,8 @@ public async Task Sampling_Stdio(string clientId)
375376            { 
376377                [ "prompt" ]  =  "Test prompt" , 
377378                [ "maxTokens" ]  =  100 
378-             } 
379-         ) ; 
379+             } , 
380+              TestContext . Current . CancellationToken ) ; 
380381
381382        // assert 
382383        Assert . NotNull ( result ) ; 
@@ -423,8 +424,8 @@ public async Task Notifications_Stdio(string clientId)
423424        await  using  var  client  =  await  _fixture . CreateClientAsync ( clientId ) ; 
424425
425426        // Verify we can send notifications without errors 
426-         await  client . SendNotificationAsync ( NotificationMethods . RootsUpdatedNotification ) ; 
427-         await  client . SendNotificationAsync ( "test/notification" ,  new  {  test  =  true  } ) ; 
427+         await  client . SendNotificationAsync ( NotificationMethods . RootsUpdatedNotification ,   cancellationToken :   TestContext . Current . CancellationToken ) ; 
428+         await  client . SendNotificationAsync ( "test/notification" ,  new  {  test  =  true  } ,   TestContext . Current . CancellationToken ) ; 
428429
429430        // assert 
430431        // no response to check, if no exception is thrown, it's a success 
@@ -452,13 +453,17 @@ public async Task CallTool_Stdio_MemoryServer()
452453            ClientInfo  =  new ( )  {  Name  =  "IntegrationTestClient" ,  Version  =  "1.0.0"  } 
453454        } ; 
454455
455-         await  using  var  client  =  await  McpClientFactory . CreateAsync ( serverConfig ,  clientOptions ,  loggerFactory :  _fixture . LoggerFactory ) ; 
456+         await  using  var  client  =  await  McpClientFactory . CreateAsync ( 
457+             serverConfig ,  
458+             clientOptions ,  
459+             loggerFactory :  _fixture . LoggerFactory ,  
460+             cancellationToken :  TestContext . Current . CancellationToken ) ; 
456461
457462        // act 
458463        var  result  =  await  client . CallToolAsync ( 
459464            "read_graph" , 
460-             [ ] 
461-         ) ; 
465+             [ ] , 
466+              TestContext . Current . CancellationToken ) ; 
462467
463468        // assert 
464469        Assert . NotNull ( result ) ; 
@@ -471,14 +476,14 @@ public async Task CallTool_Stdio_MemoryServer()
471476    [ Fact ] 
472477    public  async  Task  GetAIFunctionsAsync_UsingEverythingServer_ToolsAreProperlyCalled ( ) 
473478    { 
474-         if  ( s_openAIKey  is  null ) 
475-         { 
476-             return ;  // Skip the test if the OpenAI key is not provided 
477-         } 
479+         SkipTestIfNoOpenAIKey ( ) ; 
478480
479481        // Get the MCP client and tools from it. 
480-         await  using  var  client  =  await  McpClientFactory . CreateAsync ( _fixture . EverythingServerConfig ,  _fixture . DefaultOptions ) ;  ; 
481-         var  mappedTools  =  await  client . GetAIFunctionsAsync ( ) ; 
482+         await  using  var  client  =  await  McpClientFactory . CreateAsync ( 
483+             _fixture . EverythingServerConfig ,  
484+             _fixture . DefaultOptions ,  
485+             cancellationToken :  TestContext . Current . CancellationToken ) ; 
486+         var  mappedTools  =  await  client . GetAIFunctionsAsync ( TestContext . Current . CancellationToken ) ; 
482487
483488        // Create the chat client. 
484489        using  IChatClient  chatClient  =  new  OpenAIClient ( s_openAIKey ) . AsChatClient ( "gpt-4o-mini" ) 
@@ -495,7 +500,7 @@ public async Task GetAIFunctionsAsync_UsingEverythingServer_ToolsAreProperlyCall
495500        messages . Add ( new ( ChatRole . User ,  "Please call the echo tool with the string 'Hello MCP!' and output the response ad verbatim." ) ) ; 
496501
497502        // Call the chat client 
498-         var  response  =  await  chatClient . GetResponseAsync ( messages ,  new ( )  {  Tools  =  [ .. mappedTools ] ,  Temperature  =  0  } ) ; 
503+         var  response  =  await  chatClient . GetResponseAsync ( messages ,  new ( )  {  Tools  =  [ .. mappedTools ] ,  Temperature  =  0  } ,   TestContext . Current . CancellationToken ) ; 
499504
500505        // Assert 
501506        Assert . Contains ( "Echo: Hello MCP!" ,  response . Text ) ; 
@@ -504,10 +509,7 @@ public async Task GetAIFunctionsAsync_UsingEverythingServer_ToolsAreProperlyCall
504509    [ Fact ] 
505510    public  async  Task  SamplingViaChatClient_RequestResponseProperlyPropagated ( ) 
506511    { 
507-         if  ( s_openAIKey  is  null ) 
508-         { 
509-             return ;  // Skip the test if the OpenAI key is not provided 
510-         } 
512+         SkipTestIfNoOpenAIKey ( ) ; 
511513
512514        await  using  var  client  =  await  McpClientFactory . CreateAsync ( _fixture . EverythingServerConfig ,  new ( ) 
513515        { 
@@ -519,17 +521,22 @@ public async Task SamplingViaChatClient_RequestResponseProperlyPropagated()
519521                    SamplingHandler  =  new  OpenAIClient ( s_openAIKey ) . AsChatClient ( "gpt-4o-mini" ) . CreateSamplingHandler ( ) , 
520522                } , 
521523            } , 
522-         } ) ; 
524+         } ,   cancellationToken :   TestContext . Current . CancellationToken ) ; 
523525
524526        var  result  =  await  client . CallToolAsync ( "sampleLLM" ,  new ( ) 
525527        { 
526528            [ "prompt" ]  =  "In just a few words, what is the most famous tower in Paris?" , 
527-         } ) ; 
529+         } ,   TestContext . Current . CancellationToken ) ; 
528530
529531        Assert . NotNull ( result ) ; 
530532        Assert . NotEmpty ( result . Content ) ; 
531533        Assert . Equal ( "text" ,  result . Content [ 0 ] . Type ) ; 
532534        Assert . Contains ( "LLM sampling result:" ,  result . Content [ 0 ] . Text ) ; 
533535        Assert . Contains ( "Eiffel" ,  result . Content [ 0 ] . Text ) ; 
534536    } 
537+ 
538+     private  static void  SkipTestIfNoOpenAIKey ( ) 
539+     { 
540+         Assert . SkipWhen ( s_openAIKey  is  null ,  "No OpenAI key provided. Skipping test." ) ; 
541+     } 
535542} 
0 commit comments