mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-04-26 00:46:22 +02:00
Release/v1.2 (#457)
* Bump setup.py versions for 1.1 * PoC MCP server (#419) * Very initial MCP server PoC for TrustGraph * Put service on port 8000 * Add MCP container and packages to buildout * Update docs for API/CLI changes in 1.0 (#421) * Update some API basics for the 0.23/1.0 API change * Add MCP container push (#425) * Add command args to the MCP server (#426) * Host and port parameters * Added websocket arg * More docs * MCP client support (#427) - MCP client service - Tool request/response schema - API gateway support for mcp-tool - Message translation for tool request & response - Make mcp-tool using configuration service for information about where the MCP services are. * Feature/react call mcp (#428) Key Features - MCP Tool Integration: Added core MCP tool support with ToolClientSpec and ToolClient classes - API Enhancement: New mcp_tool method for flow-specific tool invocation - CLI Tooling: New tg-invoke-mcp-tool command for testing MCP integration - React Agent Enhancement: Fixed and improved multi-tool invocation capabilities - Tool Management: Enhanced CLI for tool configuration and management Changes - Added MCP tool invocation to API with flow-specific integration - Implemented ToolClientSpec and ToolClient for tool call handling - Updated agent-manager-react to invoke MCP tools with configurable types - Enhanced CLI with new commands and improved help text - Added comprehensive documentation for new CLI commands - Improved tool configuration management Testing - Added tg-invoke-mcp-tool CLI command for isolated MCP integration testing - Enhanced agent capability to invoke multiple tools simultaneously * Test suite executed from CI pipeline (#433) * Test strategy & test cases * Unit tests * Integration tests * Extending test coverage (#434) * Contract tests * Testing embeedings * Agent unit tests * Knowledge pipeline tests * Turn on contract tests * Increase storage test coverage (#435) * Fixing storage and adding tests * PR pipeline only runs quick tests * Empty configuration is returned as empty list, previously was not in response (#436) * Update config util to take files as well as command-line text (#437) * Updated CLI invocation and config model for tools and mcp (#438) * Updated CLI invocation and config model for tools and mcp * CLI anomalies * Tweaked the MCP tool implementation for new model * Update agent implementation to match the new model * Fix agent tools, now all tested * Fixed integration tests * Fix MCP delete tool params * Update Python deps to 1.2 * Update to enable knowledge extraction using the agent framework (#439) * Implement KG extraction agent (kg-extract-agent) * Using ReAct framework (agent-manager-react) * ReAct manager had an issue when emitting JSON, which conflicts which ReAct manager's own JSON messages, so refactored ReAct manager to use traditional ReAct messages, non-JSON structure. * Minor refactor to take the prompt template client out of prompt-template so it can be more readily used by other modules. kg-extract-agent uses this framework. * Migrate from setup.py to pyproject.toml (#440) * Converted setup.py to pyproject.toml * Modern package infrastructure as recommended by py docs * Install missing build deps (#441) * Install missing build deps (#442) * Implement logging strategy (#444) * Logging strategy and convert all prints() to logging invocations * Fix/startup failure (#445) * Fix loggin startup problems * Fix logging startup problems (#446) * Fix logging startup problems (#447) * Fixed Mistral OCR to use current API (#448) * Fixed Mistral OCR to use current API * Added PDF decoder tests * Fix Mistral OCR ident to be standard pdf-decoder (#450) * Fix Mistral OCR ident to be standard pdf-decoder * Correct test * Schema structure refactor (#451) * Write schema refactor spec * Implemented schema refactor spec * Structure data mvp (#452) * Structured data tech spec * Architecture principles * New schemas * Updated schemas and specs * Object extractor * Add .coveragerc * New tests * Cassandra object storage * Trying to object extraction working, issues exist * Validate librarian collection (#453) * Fix token chunker, broken API invocation (#454) * Fix token chunker, broken API invocation (#455) * Knowledge load utility CLI (#456) * Knowledge loader * More tests
This commit is contained in:
cybermaggedon
GitHub
parent
c85ba197be
commit
89be656990
509 changed files with 49632 additions and 5159 deletions
0
tests/unit/test_decoding/__init__.py
Normal file
0
tests/unit/test_decoding/__init__.py
Normal file
296
tests/unit/test_decoding/test_mistral_ocr_processor.py
Normal file
296
tests/unit/test_decoding/test_mistral_ocr_processor.py
Normal file
|
|
@ -0,0 +1,296 @@
|
|||
"""
|
||||
Unit tests for trustgraph.decoding.mistral_ocr.processor
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import base64
|
||||
import uuid
|
||||
from unittest.mock import AsyncMock, MagicMock, patch, Mock
|
||||
from unittest import IsolatedAsyncioTestCase
|
||||
from io import BytesIO
|
||||
|
||||
from trustgraph.decoding.mistral_ocr.processor import Processor
|
||||
from trustgraph.schema import Document, TextDocument, Metadata
|
||||
|
||||
|
||||
class TestMistralOcrProcessor(IsolatedAsyncioTestCase):
|
||||
"""Test Mistral OCR processor functionality"""
|
||||
|
||||
@patch('trustgraph.decoding.mistral_ocr.processor.Mistral')
|
||||
@patch('trustgraph.base.flow_processor.FlowProcessor.__init__')
|
||||
async def test_processor_initialization_with_api_key(self, mock_flow_init, mock_mistral_class):
|
||||
"""Test Mistral OCR processor initialization with API key"""
|
||||
# Arrange
|
||||
mock_flow_init.return_value = None
|
||||
mock_mistral = MagicMock()
|
||||
mock_mistral_class.return_value = mock_mistral
|
||||
|
||||
config = {
|
||||
'id': 'test-mistral-ocr',
|
||||
'api_key': 'test-api-key',
|
||||
'taskgroup': AsyncMock()
|
||||
}
|
||||
|
||||
# Act
|
||||
with patch.object(Processor, 'register_specification') as mock_register:
|
||||
processor = Processor(**config)
|
||||
|
||||
# Assert
|
||||
mock_flow_init.assert_called_once()
|
||||
mock_mistral_class.assert_called_once_with(api_key='test-api-key')
|
||||
|
||||
# Verify register_specification was called twice (consumer and producer)
|
||||
assert mock_register.call_count == 2
|
||||
|
||||
# Check consumer spec
|
||||
consumer_call = mock_register.call_args_list[0]
|
||||
consumer_spec = consumer_call[0][0]
|
||||
assert consumer_spec.name == "input"
|
||||
assert consumer_spec.schema == Document
|
||||
assert consumer_spec.handler == processor.on_message
|
||||
|
||||
# Check producer spec
|
||||
producer_call = mock_register.call_args_list[1]
|
||||
producer_spec = producer_call[0][0]
|
||||
assert producer_spec.name == "output"
|
||||
assert producer_spec.schema == TextDocument
|
||||
|
||||
@patch('trustgraph.base.flow_processor.FlowProcessor.__init__')
|
||||
async def test_processor_initialization_without_api_key(self, mock_flow_init):
|
||||
"""Test Mistral OCR processor initialization without API key raises error"""
|
||||
# Arrange
|
||||
mock_flow_init.return_value = None
|
||||
|
||||
config = {
|
||||
'id': 'test-mistral-ocr',
|
||||
'taskgroup': AsyncMock()
|
||||
}
|
||||
|
||||
# Act & Assert
|
||||
with patch.object(Processor, 'register_specification'):
|
||||
with pytest.raises(RuntimeError, match="Mistral API key not specified"):
|
||||
processor = Processor(**config)
|
||||
|
||||
@patch('trustgraph.decoding.mistral_ocr.processor.uuid.uuid4')
|
||||
@patch('trustgraph.decoding.mistral_ocr.processor.Mistral')
|
||||
@patch('trustgraph.base.flow_processor.FlowProcessor.__init__')
|
||||
async def test_ocr_single_chunk(self, mock_flow_init, mock_mistral_class, mock_uuid):
|
||||
"""Test OCR processing with a single chunk (less than 5 pages)"""
|
||||
# Arrange
|
||||
mock_flow_init.return_value = None
|
||||
mock_uuid.return_value = "test-uuid-1234"
|
||||
|
||||
# Mock Mistral client
|
||||
mock_mistral = MagicMock()
|
||||
mock_mistral_class.return_value = mock_mistral
|
||||
|
||||
# Mock file upload
|
||||
mock_uploaded_file = MagicMock(id="file-123")
|
||||
mock_mistral.files.upload.return_value = mock_uploaded_file
|
||||
|
||||
# Mock signed URL
|
||||
mock_signed_url = MagicMock(url="https://example.com/signed-url")
|
||||
mock_mistral.files.get_signed_url.return_value = mock_signed_url
|
||||
|
||||
# Mock OCR response
|
||||
mock_page = MagicMock(
|
||||
markdown="# Page 1\nContent ",
|
||||
images=[MagicMock(id="img1", image_base64="data:image/png;base64,abc123")]
|
||||
)
|
||||
mock_ocr_response = MagicMock(pages=[mock_page])
|
||||
mock_mistral.ocr.process.return_value = mock_ocr_response
|
||||
|
||||
# Mock PyPDF
|
||||
mock_pdf_reader = MagicMock()
|
||||
mock_pdf_reader.pages = [MagicMock(), MagicMock(), MagicMock()] # 3 pages
|
||||
|
||||
config = {
|
||||
'id': 'test-mistral-ocr',
|
||||
'api_key': 'test-api-key',
|
||||
'taskgroup': AsyncMock()
|
||||
}
|
||||
|
||||
with patch.object(Processor, 'register_specification'):
|
||||
with patch('trustgraph.decoding.mistral_ocr.processor.PdfReader', return_value=mock_pdf_reader):
|
||||
with patch('trustgraph.decoding.mistral_ocr.processor.PdfWriter') as mock_pdf_writer_class:
|
||||
mock_pdf_writer = MagicMock()
|
||||
mock_pdf_writer_class.return_value = mock_pdf_writer
|
||||
|
||||
processor = Processor(**config)
|
||||
|
||||
# Act
|
||||
result = processor.ocr(b"fake pdf content")
|
||||
|
||||
# Assert
|
||||
assert result == "# Page 1\nContent "
|
||||
|
||||
# Verify PDF writer was used to create chunk
|
||||
assert mock_pdf_writer.add_page.call_count == 3
|
||||
mock_pdf_writer.write_stream.assert_called_once()
|
||||
|
||||
# Verify Mistral API calls
|
||||
mock_mistral.files.upload.assert_called_once()
|
||||
upload_call = mock_mistral.files.upload.call_args[1]
|
||||
assert upload_call['file']['file_name'] == "test-uuid-1234"
|
||||
assert upload_call['purpose'] == 'ocr'
|
||||
|
||||
mock_mistral.files.get_signed_url.assert_called_once_with(
|
||||
file_id="file-123", expiry=1
|
||||
)
|
||||
|
||||
mock_mistral.ocr.process.assert_called_once_with(
|
||||
model="mistral-ocr-latest",
|
||||
include_image_base64=True,
|
||||
document={
|
||||
"type": "document_url",
|
||||
"document_url": "https://example.com/signed-url",
|
||||
}
|
||||
)
|
||||
|
||||
@patch('trustgraph.decoding.mistral_ocr.processor.uuid.uuid4')
|
||||
@patch('trustgraph.decoding.mistral_ocr.processor.Mistral')
|
||||
@patch('trustgraph.base.flow_processor.FlowProcessor.__init__')
|
||||
async def test_on_message_success(self, mock_flow_init, mock_mistral_class, mock_uuid):
|
||||
"""Test successful message processing"""
|
||||
# Arrange
|
||||
mock_flow_init.return_value = None
|
||||
mock_uuid.return_value = "test-uuid-5678"
|
||||
|
||||
# Mock Mistral client with simple OCR response
|
||||
mock_mistral = MagicMock()
|
||||
mock_mistral_class.return_value = mock_mistral
|
||||
|
||||
# Mock the ocr method to return simple markdown
|
||||
ocr_result = "# Document Title\nThis is the OCR content"
|
||||
|
||||
# Mock message
|
||||
pdf_content = b"fake pdf content"
|
||||
pdf_base64 = base64.b64encode(pdf_content).decode('utf-8')
|
||||
mock_metadata = Metadata(id="test-doc")
|
||||
mock_document = Document(metadata=mock_metadata, data=pdf_base64)
|
||||
mock_msg = MagicMock()
|
||||
mock_msg.value.return_value = mock_document
|
||||
|
||||
# Mock flow - needs to be a callable that returns an object with send method
|
||||
mock_output_flow = AsyncMock()
|
||||
mock_flow = MagicMock(return_value=mock_output_flow)
|
||||
|
||||
config = {
|
||||
'id': 'test-mistral-ocr',
|
||||
'api_key': 'test-api-key',
|
||||
'taskgroup': AsyncMock()
|
||||
}
|
||||
|
||||
with patch.object(Processor, 'register_specification'):
|
||||
processor = Processor(**config)
|
||||
|
||||
# Mock the ocr method
|
||||
with patch.object(processor, 'ocr', return_value=ocr_result):
|
||||
# Act
|
||||
await processor.on_message(mock_msg, None, mock_flow)
|
||||
|
||||
# Assert
|
||||
# Verify output was sent
|
||||
mock_output_flow.send.assert_called_once()
|
||||
|
||||
# Check output
|
||||
call_args = mock_output_flow.send.call_args[0][0]
|
||||
assert isinstance(call_args, TextDocument)
|
||||
assert call_args.metadata == mock_metadata
|
||||
assert call_args.text == ocr_result.encode('utf-8')
|
||||
|
||||
@patch('trustgraph.decoding.mistral_ocr.processor.Mistral')
|
||||
@patch('trustgraph.base.flow_processor.FlowProcessor.__init__')
|
||||
async def test_chunks_function(self, mock_flow_init, mock_mistral_class):
|
||||
"""Test the chunks utility function"""
|
||||
# Arrange
|
||||
from trustgraph.decoding.mistral_ocr.processor import chunks
|
||||
|
||||
test_list = list(range(12))
|
||||
|
||||
# Act
|
||||
result = list(chunks(test_list, 5))
|
||||
|
||||
# Assert
|
||||
assert len(result) == 3
|
||||
assert result[0] == [0, 1, 2, 3, 4]
|
||||
assert result[1] == [5, 6, 7, 8, 9]
|
||||
assert result[2] == [10, 11]
|
||||
|
||||
@patch('trustgraph.decoding.mistral_ocr.processor.Mistral')
|
||||
@patch('trustgraph.base.flow_processor.FlowProcessor.__init__')
|
||||
async def test_replace_images_in_markdown(self, mock_flow_init, mock_mistral_class):
|
||||
"""Test the replace_images_in_markdown function"""
|
||||
# Arrange
|
||||
from trustgraph.decoding.mistral_ocr.processor import replace_images_in_markdown
|
||||
|
||||
markdown = "# Title\n\nSome text\n"
|
||||
images_dict = {
|
||||
"image1": "data:image/png;base64,abc123",
|
||||
"image2": "data:image/png;base64,def456"
|
||||
}
|
||||
|
||||
# Act
|
||||
result = replace_images_in_markdown(markdown, images_dict)
|
||||
|
||||
# Assert
|
||||
expected = "# Title\n\nSome text\n"
|
||||
assert result == expected
|
||||
|
||||
@patch('trustgraph.decoding.mistral_ocr.processor.Mistral')
|
||||
@patch('trustgraph.base.flow_processor.FlowProcessor.__init__')
|
||||
async def test_get_combined_markdown(self, mock_flow_init, mock_mistral_class):
|
||||
"""Test the get_combined_markdown function"""
|
||||
# Arrange
|
||||
from trustgraph.decoding.mistral_ocr.processor import get_combined_markdown
|
||||
from mistralai.models import OCRResponse
|
||||
|
||||
# Mock OCR response with multiple pages
|
||||
mock_page1 = MagicMock(
|
||||
markdown="# Page 1\n",
|
||||
images=[MagicMock(id="img1", image_base64="base64_img1")]
|
||||
)
|
||||
mock_page2 = MagicMock(
|
||||
markdown="# Page 2\n",
|
||||
images=[MagicMock(id="img2", image_base64="base64_img2")]
|
||||
)
|
||||
mock_ocr_response = MagicMock(pages=[mock_page1, mock_page2])
|
||||
|
||||
# Act
|
||||
result = get_combined_markdown(mock_ocr_response)
|
||||
|
||||
# Assert
|
||||
expected = "# Page 1\n\n\n# Page 2\n"
|
||||
assert result == expected
|
||||
|
||||
@patch('trustgraph.base.flow_processor.FlowProcessor.add_args')
|
||||
def test_add_args(self, mock_parent_add_args):
|
||||
"""Test add_args adds API key argument"""
|
||||
# Arrange
|
||||
mock_parser = MagicMock()
|
||||
|
||||
# Act
|
||||
Processor.add_args(mock_parser)
|
||||
|
||||
# Assert
|
||||
mock_parent_add_args.assert_called_once_with(mock_parser)
|
||||
mock_parser.add_argument.assert_called_once_with(
|
||||
'-k', '--api-key',
|
||||
default=None, # default_api_key is None in test environment
|
||||
help='Mistral API Key'
|
||||
)
|
||||
|
||||
@patch('trustgraph.decoding.mistral_ocr.processor.Processor.launch')
|
||||
def test_run(self, mock_launch):
|
||||
"""Test run function"""
|
||||
# Act
|
||||
from trustgraph.decoding.mistral_ocr.processor import run
|
||||
run()
|
||||
|
||||
# Assert
|
||||
mock_launch.assert_called_once_with("pdf-decoder",
|
||||
"\nSimple decoder, accepts PDF documents on input, outputs pages from the\nPDF document as text as separate output objects.\n")
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
pytest.main([__file__])
|
||||
229
tests/unit/test_decoding/test_pdf_decoder.py
Normal file
229
tests/unit/test_decoding/test_pdf_decoder.py
Normal file
|
|
@ -0,0 +1,229 @@
|
|||
"""
|
||||
Unit tests for trustgraph.decoding.pdf.pdf_decoder
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import base64
|
||||
import tempfile
|
||||
from unittest.mock import AsyncMock, MagicMock, patch, call
|
||||
from unittest import IsolatedAsyncioTestCase
|
||||
|
||||
from trustgraph.decoding.pdf.pdf_decoder import Processor
|
||||
from trustgraph.schema import Document, TextDocument, Metadata
|
||||
|
||||
|
||||
class TestPdfDecoderProcessor(IsolatedAsyncioTestCase):
|
||||
"""Test PDF decoder processor functionality"""
|
||||
|
||||
@patch('trustgraph.base.flow_processor.FlowProcessor.__init__')
|
||||
async def test_processor_initialization(self, mock_flow_init):
|
||||
"""Test PDF decoder processor initialization"""
|
||||
# Arrange
|
||||
mock_flow_init.return_value = None
|
||||
|
||||
config = {
|
||||
'id': 'test-pdf-decoder',
|
||||
'taskgroup': AsyncMock()
|
||||
}
|
||||
|
||||
# Act
|
||||
with patch.object(Processor, 'register_specification') as mock_register:
|
||||
processor = Processor(**config)
|
||||
|
||||
# Assert
|
||||
mock_flow_init.assert_called_once()
|
||||
# Verify register_specification was called twice (consumer and producer)
|
||||
assert mock_register.call_count == 2
|
||||
|
||||
# Check consumer spec
|
||||
consumer_call = mock_register.call_args_list[0]
|
||||
consumer_spec = consumer_call[0][0]
|
||||
assert consumer_spec.name == "input"
|
||||
assert consumer_spec.schema == Document
|
||||
assert consumer_spec.handler == processor.on_message
|
||||
|
||||
# Check producer spec
|
||||
producer_call = mock_register.call_args_list[1]
|
||||
producer_spec = producer_call[0][0]
|
||||
assert producer_spec.name == "output"
|
||||
assert producer_spec.schema == TextDocument
|
||||
|
||||
@patch('trustgraph.decoding.pdf.pdf_decoder.PyPDFLoader')
|
||||
@patch('trustgraph.base.flow_processor.FlowProcessor.__init__')
|
||||
async def test_on_message_success(self, mock_flow_init, mock_pdf_loader_class):
|
||||
"""Test successful PDF processing"""
|
||||
# Arrange
|
||||
mock_flow_init.return_value = None
|
||||
|
||||
# Mock PDF content
|
||||
pdf_content = b"fake pdf content"
|
||||
pdf_base64 = base64.b64encode(pdf_content).decode('utf-8')
|
||||
|
||||
# Mock PyPDFLoader
|
||||
mock_loader = MagicMock()
|
||||
mock_page1 = MagicMock(page_content="Page 1 content")
|
||||
mock_page2 = MagicMock(page_content="Page 2 content")
|
||||
mock_loader.load.return_value = [mock_page1, mock_page2]
|
||||
mock_pdf_loader_class.return_value = mock_loader
|
||||
|
||||
# Mock message
|
||||
mock_metadata = Metadata(id="test-doc")
|
||||
mock_document = Document(metadata=mock_metadata, data=pdf_base64)
|
||||
mock_msg = MagicMock()
|
||||
mock_msg.value.return_value = mock_document
|
||||
|
||||
# Mock flow - needs to be a callable that returns an object with send method
|
||||
mock_output_flow = AsyncMock()
|
||||
mock_flow = MagicMock(return_value=mock_output_flow)
|
||||
|
||||
config = {
|
||||
'id': 'test-pdf-decoder',
|
||||
'taskgroup': AsyncMock()
|
||||
}
|
||||
|
||||
with patch.object(Processor, 'register_specification'):
|
||||
processor = Processor(**config)
|
||||
|
||||
# Act
|
||||
await processor.on_message(mock_msg, None, mock_flow)
|
||||
|
||||
# Assert
|
||||
# Verify PyPDFLoader was called
|
||||
mock_pdf_loader_class.assert_called_once()
|
||||
mock_loader.load.assert_called_once()
|
||||
|
||||
# Verify output was sent for each page
|
||||
assert mock_output_flow.send.call_count == 2
|
||||
|
||||
# Check first page output
|
||||
first_call = mock_output_flow.send.call_args_list[0]
|
||||
first_output = first_call[0][0]
|
||||
assert isinstance(first_output, TextDocument)
|
||||
assert first_output.metadata == mock_metadata
|
||||
assert first_output.text == b"Page 1 content"
|
||||
|
||||
# Check second page output
|
||||
second_call = mock_output_flow.send.call_args_list[1]
|
||||
second_output = second_call[0][0]
|
||||
assert isinstance(second_output, TextDocument)
|
||||
assert second_output.metadata == mock_metadata
|
||||
assert second_output.text == b"Page 2 content"
|
||||
|
||||
@patch('trustgraph.decoding.pdf.pdf_decoder.PyPDFLoader')
|
||||
@patch('trustgraph.base.flow_processor.FlowProcessor.__init__')
|
||||
async def test_on_message_empty_pdf(self, mock_flow_init, mock_pdf_loader_class):
|
||||
"""Test handling of empty PDF"""
|
||||
# Arrange
|
||||
mock_flow_init.return_value = None
|
||||
|
||||
# Mock PDF content
|
||||
pdf_content = b"fake pdf content"
|
||||
pdf_base64 = base64.b64encode(pdf_content).decode('utf-8')
|
||||
|
||||
# Mock PyPDFLoader with no pages
|
||||
mock_loader = MagicMock()
|
||||
mock_loader.load.return_value = []
|
||||
mock_pdf_loader_class.return_value = mock_loader
|
||||
|
||||
# Mock message
|
||||
mock_metadata = Metadata(id="test-doc")
|
||||
mock_document = Document(metadata=mock_metadata, data=pdf_base64)
|
||||
mock_msg = MagicMock()
|
||||
mock_msg.value.return_value = mock_document
|
||||
|
||||
# Mock flow - needs to be a callable that returns an object with send method
|
||||
mock_output_flow = AsyncMock()
|
||||
mock_flow = MagicMock(return_value=mock_output_flow)
|
||||
|
||||
config = {
|
||||
'id': 'test-pdf-decoder',
|
||||
'taskgroup': AsyncMock()
|
||||
}
|
||||
|
||||
with patch.object(Processor, 'register_specification'):
|
||||
processor = Processor(**config)
|
||||
|
||||
# Act
|
||||
await processor.on_message(mock_msg, None, mock_flow)
|
||||
|
||||
# Assert
|
||||
# Verify PyPDFLoader was called
|
||||
mock_pdf_loader_class.assert_called_once()
|
||||
mock_loader.load.assert_called_once()
|
||||
|
||||
# Verify no output was sent
|
||||
mock_output_flow.send.assert_not_called()
|
||||
|
||||
@patch('trustgraph.decoding.pdf.pdf_decoder.PyPDFLoader')
|
||||
@patch('trustgraph.base.flow_processor.FlowProcessor.__init__')
|
||||
async def test_on_message_unicode_content(self, mock_flow_init, mock_pdf_loader_class):
|
||||
"""Test handling of unicode content in PDF"""
|
||||
# Arrange
|
||||
mock_flow_init.return_value = None
|
||||
|
||||
# Mock PDF content
|
||||
pdf_content = b"fake pdf content"
|
||||
pdf_base64 = base64.b64encode(pdf_content).decode('utf-8')
|
||||
|
||||
# Mock PyPDFLoader with unicode content
|
||||
mock_loader = MagicMock()
|
||||
mock_page = MagicMock(page_content="Page with unicode: 你好世界 🌍")
|
||||
mock_loader.load.return_value = [mock_page]
|
||||
mock_pdf_loader_class.return_value = mock_loader
|
||||
|
||||
# Mock message
|
||||
mock_metadata = Metadata(id="test-doc")
|
||||
mock_document = Document(metadata=mock_metadata, data=pdf_base64)
|
||||
mock_msg = MagicMock()
|
||||
mock_msg.value.return_value = mock_document
|
||||
|
||||
# Mock flow - needs to be a callable that returns an object with send method
|
||||
mock_output_flow = AsyncMock()
|
||||
mock_flow = MagicMock(return_value=mock_output_flow)
|
||||
|
||||
config = {
|
||||
'id': 'test-pdf-decoder',
|
||||
'taskgroup': AsyncMock()
|
||||
}
|
||||
|
||||
with patch.object(Processor, 'register_specification'):
|
||||
processor = Processor(**config)
|
||||
|
||||
# Act
|
||||
await processor.on_message(mock_msg, None, mock_flow)
|
||||
|
||||
# Assert
|
||||
# Verify output was sent
|
||||
mock_output_flow.send.assert_called_once()
|
||||
|
||||
# Check output
|
||||
call_args = mock_output_flow.send.call_args[0][0]
|
||||
assert isinstance(call_args, TextDocument)
|
||||
assert call_args.text == "Page with unicode: 你好世界 🌍".encode('utf-8')
|
||||
|
||||
@patch('trustgraph.base.flow_processor.FlowProcessor.add_args')
|
||||
def test_add_args(self, mock_parent_add_args):
|
||||
"""Test add_args calls parent method"""
|
||||
# Arrange
|
||||
mock_parser = MagicMock()
|
||||
|
||||
# Act
|
||||
Processor.add_args(mock_parser)
|
||||
|
||||
# Assert
|
||||
mock_parent_add_args.assert_called_once_with(mock_parser)
|
||||
|
||||
@patch('trustgraph.decoding.pdf.pdf_decoder.Processor.launch')
|
||||
def test_run(self, mock_launch):
|
||||
"""Test run function"""
|
||||
# Act
|
||||
from trustgraph.decoding.pdf.pdf_decoder import run
|
||||
run()
|
||||
|
||||
# Assert
|
||||
mock_launch.assert_called_once_with("pdf-decoder",
|
||||
"\nSimple decoder, accepts PDF documents on input, outputs pages from the\nPDF document as text as separate output objects.\n")
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
pytest.main([__file__])
|
||||
Loading…
Add table
Add a link
Reference in a new issue