Mocked db_recorder object in tests

Author: aponcedeleonch

src/codegate/pipeline/output.py

@@ -76,6 +76,7 @@ def __init__(
         self,
         pipeline_steps: list[OutputPipelineStep],
         input_context: Optional[PipelineContext] = None,
+        db_recorder: Optional[DbRecorder] = None,
     ):
         self._input_context = input_context
         self._pipeline_steps = pipeline_steps
@@ -84,7 +85,10 @@ def __init__(
         # the remaining content in the buffer when the stream ends, we need
         # to store the parameters like model, timestamp, etc.
         self._buffered_chunk = None
-        self._db_recorder = DbRecorder()
+        if not db_recorder:
+            self._db_recorder = DbRecorder()
+        else:
+            self._db_recorder = db_recorder
 
     def _buffer_chunk(self, chunk: ModelResponse) -> None:
         """

tests/pipeline/test_output.py

@@ -1,4 +1,5 @@
 from typing import List
+from unittest.mock import MagicMock
 
 import pytest
 from litellm import ModelResponse
@@ -94,7 +95,8 @@ async def test_single_step_processing(self):
         """Test processing a stream through a single step"""
         step = MockOutputPipelineStep("test_step", modify_content=True)
         context = MockContext()
-        instance = OutputPipelineInstance([step], context)
+        db_recorder = MagicMock()
+        instance = OutputPipelineInstance([step], context, db_recorder)
 
         async def mock_stream():
             yield create_model_response("Hello")
@@ -118,7 +120,8 @@ async def test_multiple_steps_processing(self):
             MockOutputPipelineStep("step2", modify_content=True),
         ]
         context = MockContext()
-        instance = OutputPipelineInstance(steps, context)
+        db_recorder = MagicMock()
+        instance = OutputPipelineInstance(steps, context, db_recorder)
 
         async def mock_stream():
             yield create_model_response("Hello")
@@ -197,7 +200,8 @@ async def process_chunk(
                 return []
 
         context = MockContext()
-        instance = OutputPipelineInstance([ReplacementStep()], context)
+        db_recorder = MagicMock()
+        instance = OutputPipelineInstance([ReplacementStep()], context, db_recorder)
 
         async def mock_stream():
             yield create_model_response("he")
@@ -221,7 +225,8 @@ async def test_buffer_processing(self):
         """Test that content is properly buffered and cleared"""
         step = MockOutputPipelineStep("test_step")
         context = MockContext()
-        instance = OutputPipelineInstance([step], context)
+        db_recorder = MagicMock()
+        instance = OutputPipelineInstance([step], context, db_recorder)
 
         async def mock_stream():
             yield create_model_response("Hello")
@@ -242,7 +247,8 @@ async def test_empty_stream(self):
         """Test handling of an empty stream"""
         step = MockOutputPipelineStep("test_step")
         context = MockContext()
-        instance = OutputPipelineInstance([step], context)
+        db_recorder = MagicMock()
+        instance = OutputPipelineInstance([step], context, db_recorder)
 
         async def mock_stream():
             if False:
@@ -275,7 +281,10 @@ async def process_chunk(
                 assert input_context.metadata["test"] == "value"
                 return [chunk]
 
-        instance = OutputPipelineInstance([ContextCheckingStep()], input_context=input_context)
+        db_recorder = MagicMock()
+        instance = OutputPipelineInstance(
+            [ContextCheckingStep()], input_context=input_context, db_recorder=db_recorder
+        )
 
         async def mock_stream():
             yield create_model_response("test")
@@ -288,7 +297,8 @@ async def test_buffer_flush_on_stream_end(self):
         """Test that buffer is properly flushed when stream ends"""
         step = MockOutputPipelineStep("test_step", should_pause=True)
         context = MockContext()
-        instance = OutputPipelineInstance([step], context)
+        db_recorder = MagicMock()
+        instance = OutputPipelineInstance([step], context, db_recorder)
 
         async def mock_stream():
             yield create_model_response("Hello")