From fc5ec0f72a7a698b90b73155586703d922f99b0f Mon Sep 17 00:00:00 2001
From: Javier Martinez <javiermartinezalvarez98@gmail.com>
Date: Tue, 14 Jul 2026 17:03:45 +0200
Subject: [PATCH] fix: tests

---
 tests/server/chat/test_chat_routes.py              | 14 ++++++++++++--
 .../chat/test_chat_routes_skills_integration.py    |  6 +++---
 2 files changed, 15 insertions(+), 5 deletions(-)

diff --git a/tests/server/chat/test_chat_routes.py b/tests/server/chat/test_chat_routes.py
index a8e488c0..b1884f9e 100644
--- a/tests/server/chat/test_chat_routes.py
+++ b/tests/server/chat/test_chat_routes.py
@@ -2264,14 +2264,24 @@ async def test_principal_is_propagated_to_background_task(
     original_process_stream = stream_processor.process_stream
 
     async def capturing_process_stream(
-        correlation_id, stream_type, event_generator, event_handler, metadata=None
+        correlation_id,
+        stream_type,
+        event_generator,
+        event_handler,
+        metadata=None,
+        mark_completed=True,
     ):
         # This coroutine runs inside the asyncio task spawned by TaskManager.
         # The task was created with copy_context(), so Principal.current() here
         # must return the principal that was active during the HTTP request.
         captured.append(Principal.current().api_key)
         await original_process_stream(
-            correlation_id, stream_type, event_generator, event_handler, metadata
+            correlation_id,
+            stream_type,
+            event_generator,
+            event_handler,
+            metadata,
+            mark_completed,
         )
 
     stream_processor.process_stream = capturing_process_stream  # type: ignore[method-assign]
diff --git a/tests/server/chat/test_chat_routes_skills_integration.py b/tests/server/chat/test_chat_routes_skills_integration.py
index 4b98ad92..488272d2 100644
--- a/tests/server/chat/test_chat_routes_skills_integration.py
+++ b/tests/server/chat/test_chat_routes_skills_integration.py
@@ -1,6 +1,7 @@
 import json
 import uuid
 from typing import Any
+from unittest.mock import Mock
 
 import pytest
 from httpx import AsyncClient
@@ -72,9 +73,8 @@ async def mock_llm_with_capture(
 
     mock_llm.astream_chat_with_tools = coro
     llm_component = injector.get(LLMComponent)
-    llm_component.llm = mock_llm
-    llm_component.get_llm.side_effect = lambda *args, **kwargs: mock_llm
-    injector.bind_mock(LLMComponent, mock_llm)
+    llm_component.get_llm = Mock(return_value=mock_llm)
+    injector.bind_mock(LLMComponent, llm_component)
 
 
 def _skill_md(name: str, description: str, body: str) -> bytes: