@@ -1446,10 +1446,14 @@ func TestSSEServer(t *testing.T) {
1446
1446
1447
1447
t .Run ("Headers are passed through to tool requests" , func (t * testing.T ) {
1448
1448
hooks := & Hooks {}
1449
- var headerVerified bool
1449
+ headerVerified := make ( chan struct {})
1450
1450
hooks .AddAfterCallTool (func (ctx context.Context , id any , message * mcp.CallToolRequest , result * mcp.CallToolResult ) {
1451
1451
if message .Params .Name == "verify-headers" {
1452
- headerVerified = true
1452
+ select {
1453
+ case <- headerVerified :
1454
+ default :
1455
+ close (headerVerified )
1456
+ }
1453
1457
}
1454
1458
})
1455
1459
@@ -1506,20 +1510,24 @@ func TestSSEServer(t *testing.T) {
1506
1510
t .Errorf ("Expected status 202, got %d" , resp .StatusCode )
1507
1511
}
1508
1512
1509
- // Wait briefly for hook to be called
1510
- time . Sleep ( 100 * time . Millisecond )
1511
-
1512
- if ! headerVerified {
1513
- t .Error ("Header verification hook was not called" )
1513
+ // Wait for hook to be called
1514
+ select {
1515
+ case <- headerVerified :
1516
+ case <- time . After ( 1 * time . Second ):
1517
+ t .Error ("Header verification hook was not called within timeout " )
1514
1518
}
1515
1519
})
1516
1520
1517
1521
t .Run ("Headers are not nil when no headers are set" , func (t * testing.T ) {
1518
1522
hooks := & Hooks {}
1519
- var headersChecked bool
1523
+ headersChecked := make ( chan struct {})
1520
1524
hooks .AddAfterCallTool (func (ctx context.Context , id any , message * mcp.CallToolRequest , result * mcp.CallToolResult ) {
1521
1525
if message .Params .Name == "check-headers-not-nil" {
1522
- headersChecked = true
1526
+ select {
1527
+ case <- headersChecked :
1528
+ default :
1529
+ close (headersChecked )
1530
+ }
1523
1531
}
1524
1532
})
1525
1533
@@ -1591,11 +1599,11 @@ func TestSSEServer(t *testing.T) {
1591
1599
t .Errorf ("Expected status 202, got %d" , resp .StatusCode )
1592
1600
}
1593
1601
1594
- // Wait briefly for hook to be called
1595
- time . Sleep ( 100 * time . Millisecond )
1596
-
1597
- if ! headersChecked {
1598
- t .Error ("Headers check hook was not called" )
1602
+ // Wait for hook to be called
1603
+ select {
1604
+ case <- headersChecked :
1605
+ case <- time . After ( 1 * time . Second ):
1606
+ t .Error ("Headers check hook was not called within timeout " )
1599
1607
}
1600
1608
})
1601
1609
}
0 commit comments