WIP
This commit is contained in:
@@ -232,10 +232,8 @@ class TestInferenceServicePDFRendering:
|
||||
@patch('backend.pipeline.pipeline.InferencePipeline')
|
||||
@patch('backend.pipeline.yolo_detector.YOLODetector')
|
||||
@patch('shared.pdf.renderer.render_pdf_to_images')
|
||||
@patch('ultralytics.YOLO')
|
||||
def test_pdf_visualization_imports_correctly(
|
||||
self,
|
||||
mock_yolo_class,
|
||||
mock_render_pdf,
|
||||
mock_yolo_detector,
|
||||
mock_pipeline,
|
||||
@@ -248,12 +246,22 @@ class TestInferenceServicePDFRendering:
|
||||
This catches the import error we had with:
|
||||
from ..pdf.renderer (wrong) vs from shared.pdf.renderer (correct)
|
||||
"""
|
||||
# Setup mocks
|
||||
# Setup mocks for detector
|
||||
mock_detector_instance = Mock()
|
||||
mock_pipeline_instance = Mock()
|
||||
mock_model = Mock()
|
||||
mock_result = Mock()
|
||||
mock_result.save = Mock()
|
||||
mock_model.predict.return_value = [mock_result]
|
||||
mock_detector_instance.model = mock_model
|
||||
mock_yolo_detector.return_value = mock_detector_instance
|
||||
|
||||
# Setup mock for pipeline
|
||||
mock_pipeline_instance = Mock()
|
||||
mock_pipeline.return_value = mock_pipeline_instance
|
||||
|
||||
# Initialize service to setup _detector
|
||||
inference_service.initialize()
|
||||
|
||||
# Create a fake PDF path
|
||||
pdf_path = tmp_path / "test.pdf"
|
||||
pdf_path.touch()
|
||||
@@ -264,18 +272,12 @@ class TestInferenceServicePDFRendering:
|
||||
img.save(image_bytes, format='PNG')
|
||||
mock_render_pdf.return_value = [(1, image_bytes.getvalue())]
|
||||
|
||||
# Mock YOLO
|
||||
mock_model_instance = Mock()
|
||||
mock_result = Mock()
|
||||
mock_result.save = Mock()
|
||||
mock_model_instance.predict.return_value = [mock_result]
|
||||
mock_yolo_class.return_value = mock_model_instance
|
||||
|
||||
# This should not raise ImportError
|
||||
# This should not raise ImportError and should use self._detector.model
|
||||
result_path = inference_service._save_pdf_visualization(pdf_path, "test123")
|
||||
|
||||
# Verify import was successful
|
||||
# Verify import was successful and detector.model was used
|
||||
mock_render_pdf.assert_called_once()
|
||||
mock_model.predict.assert_called_once()
|
||||
assert result_path is not None
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user