Fixing tests

This commit is contained in:
Cyber MacGeddon 2025-09-26 09:22:04 +01:00
parent 5be9f563a0
commit 51d6319a5e
2 changed files with 36 additions and 25 deletions

View file

@ -9,7 +9,7 @@ from unittest import IsolatedAsyncioTestCase
# Import the service under test # Import the service under test
from trustgraph.chunking.recursive.chunker import Processor from trustgraph.chunking.recursive.chunker import Processor
from trustgraph.schema import TextDocument, Chunk from trustgraph.schema import TextDocument, Chunk, Metadata
class MockAsyncProcessor: class MockAsyncProcessor:
@ -167,32 +167,35 @@ class TestRecursiveChunkerSimple(IsolatedAsyncioTestCase):
# Mock message with TextDocument # Mock message with TextDocument
mock_message = MagicMock() mock_message = MagicMock()
mock_text_doc = MagicMock() mock_text_doc = MagicMock()
mock_text_doc.metadata.id = "test-doc-123" mock_text_doc.metadata = Metadata(
id="test-doc-123",
metadata=[],
user="test-user",
collection="test-collection"
)
mock_text_doc.text = b"This is test document content" mock_text_doc.text = b"This is test document content"
mock_message.value.return_value = mock_text_doc mock_message.value.return_value = mock_text_doc
# Mock consumer and flow with parameter overrides # Mock consumer and flow with parameter overrides
mock_consumer = MagicMock() mock_consumer = MagicMock()
mock_producer = AsyncMock()
mock_flow = MagicMock() mock_flow = MagicMock()
mock_flow.side_effect = lambda param: { mock_flow.side_effect = lambda param: {
"chunk-size": 1500, "chunk-size": 1500,
"chunk-overlap": 150 "chunk-overlap": 150,
"output": mock_producer
}.get(param) }.get(param)
mock_producer = AsyncMock()
mock_flow.return_value = mock_producer
# Act # Act
await processor.on_message(mock_message, mock_consumer, mock_flow) await processor.on_message(mock_message, mock_consumer, mock_flow)
# Assert # Assert
# Verify RecursiveCharacterTextSplitter was created with overridden parameters # Verify RecursiveCharacterTextSplitter was called with overridden parameters (last call)
mock_splitter_class.assert_called_once_with( actual_last_call = mock_splitter_class.call_args_list[-1]
chunk_size=1500, assert actual_last_call.kwargs['chunk_size'] == 1500
chunk_overlap=150, assert actual_last_call.kwargs['chunk_overlap'] == 150
length_function=len, assert actual_last_call.kwargs['length_function'] == len
is_separator_regex=False, assert actual_last_call.kwargs['is_separator_regex'] == False
)
# Verify chunk was sent to output # Verify chunk was sent to output
mock_producer.send.assert_called_once() mock_producer.send.assert_called_once()

View file

@ -9,7 +9,7 @@ from unittest import IsolatedAsyncioTestCase
# Import the service under test # Import the service under test
from trustgraph.chunking.token.chunker import Processor from trustgraph.chunking.token.chunker import Processor
from trustgraph.schema import TextDocument, Chunk from trustgraph.schema import TextDocument, Chunk, Metadata
class MockAsyncProcessor: class MockAsyncProcessor:
@ -167,31 +167,39 @@ class TestTokenChunkerSimple(IsolatedAsyncioTestCase):
# Mock message with TextDocument # Mock message with TextDocument
mock_message = MagicMock() mock_message = MagicMock()
mock_text_doc = MagicMock() mock_text_doc = MagicMock()
mock_text_doc.metadata.id = "test-doc-456" mock_text_doc.metadata = Metadata(
id="test-doc-456",
metadata=[],
user="test-user",
collection="test-collection"
)
mock_text_doc.text = b"This is test document content for token chunking" mock_text_doc.text = b"This is test document content for token chunking"
mock_message.value.return_value = mock_text_doc mock_message.value.return_value = mock_text_doc
# Mock consumer and flow with parameter overrides # Mock consumer and flow with parameter overrides
mock_consumer = MagicMock() mock_consumer = MagicMock()
mock_producer = AsyncMock()
mock_flow = MagicMock() mock_flow = MagicMock()
mock_flow.side_effect = lambda param: { mock_flow.side_effect = lambda param: {
"chunk-size": 400, "chunk-size": 400,
"chunk-overlap": 40 "chunk-overlap": 40,
"output": mock_producer
}.get(param) }.get(param)
mock_producer = AsyncMock()
mock_flow.return_value = mock_producer
# Act # Act
await processor.on_message(mock_message, mock_consumer, mock_flow) await processor.on_message(mock_message, mock_consumer, mock_flow)
# Assert # Assert
# Verify TokenTextSplitter was created with overridden parameters # Verify TokenTextSplitter was called with overridden parameters (last call)
mock_splitter_class.assert_called_once_with( expected_call = [
encoding_name="cl100k_base", ('encoding_name', 'cl100k_base'),
chunk_size=400, ('chunk_size', 400),
chunk_overlap=40, ('chunk_overlap', 40)
) ]
actual_last_call = mock_splitter_class.call_args_list[-1]
assert actual_last_call.kwargs['encoding_name'] == "cl100k_base"
assert actual_last_call.kwargs['chunk_size'] == 400
assert actual_last_call.kwargs['chunk_overlap'] == 40
# Verify chunk was sent to output # Verify chunk was sent to output
mock_producer.send.assert_called_once() mock_producer.send.assert_called_once()