mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-07-24 04:31:02 +02:00
Fixing tests
This commit is contained in:
parent
5be9f563a0
commit
51d6319a5e
2 changed files with 36 additions and 25 deletions
|
|
@ -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()
|
||||||
|
|
|
||||||
|
|
@ -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()
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue