apunkt/trustgraph
1
0
Fork 0
mirror of https://github.com/trustgraph-ai/trustgraph.git synced 2026-07-23 20:21:03 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions 3

Fixing tests

Browse source
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
Download patch file Download diff file Expand all files Collapse all files

29
tests/unit/test_chunking/test_recursive_chunker.py
View file

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

32
tests/unit/test_chunking/test_token_chunker.py
View file

@ -9,7 +9,7 @@ from unittest import IsolatedAsyncioTestCase
# Import the service under test
from trustgraph.chunking.token.chunker import Processor
from trustgraph.schema import TextDocument, Chunk
from trustgraph.schema import TextDocument, Chunk, Metadata
class MockAsyncProcessor:
@ -167,31 +167,39 @@ class TestTokenChunkerSimple(IsolatedAsyncioTestCase):
# Mock message with TextDocument
mock_message = 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_message.value.return_value = mock_text_doc
# Mock consumer and flow with parameter overrides
mock_consumer = MagicMock()
mock_producer = AsyncMock()
mock_flow = MagicMock()
mock_flow.side_effect = lambda param: {
"chunk-size": 400,
"chunk-overlap": 40
"chunk-overlap": 40,
"output": mock_producer
}.get(param)
mock_producer = AsyncMock()
mock_flow.return_value = mock_producer
# Act
await processor.on_message(mock_message, mock_consumer, mock_flow)
# Assert
# Verify TokenTextSplitter was created with overridden parameters
mock_splitter_class.assert_called_once_with(
encoding_name="cl100k_base",
chunk_size=400,
chunk_overlap=40,
)
# Verify TokenTextSplitter was called with overridden parameters (last call)
expected_call = [
('encoding_name', 'cl100k_base'),
('chunk_size', 400),
('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
mock_producer.send.assert_called_once()