Fix RetrievalWorkflow unit tests: Pass mock Context to step methods

Author: gabaoun

tests/test_retrieval.py

@@ -7,23 +7,26 @@
 async def test_retrieval_cache_hit(mocker):
     mock_chroma = MagicMock()
     wf = RetrievalWorkflow(chroma_service=mock_chroma, reranker=None)
+    mock_ctx = MagicMock()
 
     # Mock cache hit
     mocker.patch.object(wf.cache, 'get_cache', return_value="Cached Answer")
 
     # Execution
     ev = StartEvent(query="What is Aether?")
-    result = await wf.process_start(ev)
+    result = await wf.process_start(mock_ctx, ev)
 
     # Validation
     assert isinstance(result, StopEvent)
     assert result.result["answer"] == "Cached Answer"
     assert result.result["from_cache"] is True
+    mock_ctx.send_event.assert_called()
 
 @pytest.mark.asyncio
 async def test_retrieval_query_transformation(mocker):
     mock_chroma = MagicMock()
     wf = RetrievalWorkflow(chroma_service=mock_chroma, reranker=None)
+    mock_ctx = MagicMock()
 
     # Mock cache miss
     mocker.patch.object(wf.cache, 'get_cache', return_value=None)
@@ -35,9 +38,10 @@ async def test_retrieval_query_transformation(mocker):
 
     # Execution
     ev = StartEvent(query="New Query")
-    result = await wf.process_start(ev)
+    result = await wf.process_start(mock_ctx, ev)
 
     # Validation
     assert isinstance(result, QueryTransformedEvent)
     assert result.query_bundle.query_str == "New Query"
     assert "Hypothetical document content" in result.query_bundle.custom_embedding_strs
+    mock_ctx.send_event.assert_called()