diff --git a/tests/Feature/SearchAndSummarizeTest.php b/tests/Feature/SearchAndSummarizeTest.php
index a67cec3f..7bd8c4a1 100644
--- a/tests/Feature/SearchAndSummarizeTest.php
+++ b/tests/Feature/SearchAndSummarizeTest.php
@@ -61,7 +61,7 @@ public function test_gets_user_input()
             ]),
         ]);
 
-        LlmDriverFacade::shouldReceive('driver->chat')
+        LlmDriverFacade::shouldReceive('driver->setToolType->chat')
             ->once()
             ->andReturn($dto);