Fixing tests

This commit is contained in:
Cyber MacGeddon 2026-03-09 10:08:44 +00:00
parent 4e3db11323
commit 9a4bfafb25
6 changed files with 33 additions and 14 deletions

View file

@ -8,6 +8,7 @@ response delivery through the complete pipeline.
import pytest
from unittest.mock import AsyncMock
from trustgraph.retrieval.document_rag.document_rag import DocumentRag
from trustgraph.schema import ChunkMatch
from tests.utils.streaming_assertions import (
assert_streaming_chunks_valid,
assert_callback_invoked,
@ -36,10 +37,14 @@ class TestDocumentRagStreaming:
@pytest.fixture
def mock_doc_embeddings_client(self):
"""Mock document embeddings client that returns chunk IDs"""
"""Mock document embeddings client that returns chunk matches"""
client = AsyncMock()
# Now returns chunk_ids instead of actual content
client.query.return_value = ["doc/c1", "doc/c2", "doc/c3"]
# Returns ChunkMatch objects with chunk_id and score
client.query.return_value = [
ChunkMatch(chunk_id="doc/c1", score=0.95),
ChunkMatch(chunk_id="doc/c2", score=0.90),
ChunkMatch(chunk_id="doc/c3", score=0.85)
]
return client
@pytest.fixture