fix: tests

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]

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: