diff --git a/tests/test_headless_service.py b/tests/test_headless_service.py
index 066f2086..4092efe8 100644
--- a/tests/test_headless_service.py
+++ b/tests/test_headless_service.py
@@ -8,6 +8,7 @@ from fastapi.testclient import TestClient
 sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
 
 from src.app_controller import AppController
+from src.result_types import Result
 
 class TestHeadlessAPI(unittest.TestCase):
     def setUp(self) -> None:
@@ -55,7 +56,7 @@ class TestHeadlessAPI(unittest.TestCase):
      self.assertIn("not configured", response.json()["detail"])
 
     def test_generate_endpoint(self) -> None:
-     with patch('src.ai_client.send', return_value="AI Response"), \
+     with patch('src.ai_client.send_result', return_value=Result(data="AI Response")), \
           patch('src.app_controller.AppController._do_generate', return_value=("md", "path", [], "stable", "disc")):
       payload = {"prompt": "test prompt", "auto_add_history": False}
       response = self.client.post("/api/v1/generate", json=payload, headers=self.headers)