WIP

tests/inference/test_pipeline.py

@@ -497,5 +497,178 @@ class TestExtractBusinessFeaturesErrorHandling:
             assert "NumericException" in result.errors[0]
 
 
+class TestProcessPdfTokenPath:
+    """Tests for PDF text token extraction path in process_pdf()."""
+
+    def _make_pipeline(self):
+        """Create pipeline with mocked internals, bypassing __init__."""
+        with patch.object(InferencePipeline, '__init__', lambda self, **kw: None):
+            p = InferencePipeline()
+        p.detector = MagicMock()
+        p.extractor = MagicMock()
+        p.payment_line_parser = MagicMock()
+        p.dpi = 300
+        p.enable_fallback = False
+        p.enable_business_features = False
+        p.vat_tolerance = 0.5
+        p.line_items_extractor = None
+        p.vat_extractor = None
+        p.vat_validator = None
+        p._business_ocr_engine = None
+        p._table_detector = None
+        return p
+
+    def _make_detection(self, class_name='Amount', confidence=0.85, page_no=0):
+        """Create a Detection object."""
+        from backend.pipeline.yolo_detector import Detection
+        return Detection(
+            class_id=6,
+            class_name=class_name,
+            confidence=confidence,
+            bbox=(100.0, 200.0, 300.0, 250.0),
+            page_no=page_no,
+        )
+
+    def _make_extracted_field(self, field_name='Amount', raw_text='2.254,50',
+                              normalized='2254.50', confidence=0.85):
+        """Create an ExtractedField object."""
+        from backend.pipeline.field_extractor import ExtractedField
+        return ExtractedField(
+            field_name=field_name,
+            raw_text=raw_text,
+            normalized_value=normalized,
+            confidence=confidence,
+            detection_confidence=confidence,
+            ocr_confidence=1.0,
+            bbox=(100.0, 200.0, 300.0, 250.0),
+            page_no=0,
+        )
+
+    def _make_image_bytes(self):
+        """Create minimal valid PNG bytes (100x100 white image)."""
+        from PIL import Image as PILImage
+        import io as _io
+        img = PILImage.new('RGB', (100, 100), color='white')
+        buf = _io.BytesIO()
+        img.save(buf, format='PNG')
+        return buf.getvalue()
+
+    @patch('shared.pdf.extractor.PDFDocument')
+    @patch('shared.pdf.renderer.render_pdf_to_images')
+    def test_text_pdf_uses_pdf_tokens(self, mock_render, mock_pdf_doc_cls):
+        """When PDF is text-based, extract_from_detection_with_pdf is used."""
+        from shared.pdf.extractor import Token
+
+        pipeline = self._make_pipeline()
+        detection = self._make_detection()
+        image_bytes = self._make_image_bytes()
+
+        # Setup PDFDocument mock - text PDF with tokens
+        mock_pdf_doc = MagicMock()
+        mock_pdf_doc.is_text_pdf.return_value = True
+        mock_pdf_doc.page_count = 1
+        tokens = [Token(text="2.254,50", bbox=(100, 200, 200, 220), page_no=0)]
+        mock_pdf_doc.extract_text_tokens.return_value = iter(tokens)
+        mock_pdf_doc_cls.return_value.__enter__ = MagicMock(return_value=mock_pdf_doc)
+        mock_pdf_doc_cls.return_value.__exit__ = MagicMock(return_value=False)
+
+        pipeline.detector.detect.return_value = [detection]
+        pipeline.extractor.extract_from_detection_with_pdf.return_value = (
+            self._make_extracted_field()
+        )
+
+        mock_render.return_value = iter([(0, image_bytes)])
+        result = pipeline.process_pdf('/fake/invoice.pdf')
+
+        pipeline.extractor.extract_from_detection_with_pdf.assert_called_once()
+        pipeline.extractor.extract_from_detection.assert_not_called()
+        assert result.fields.get('Amount') == '2254.50'
+        assert result.success is True
+
+    @patch('shared.pdf.extractor.PDFDocument')
+    @patch('shared.pdf.renderer.render_pdf_to_images')
+    def test_scanned_pdf_uses_ocr(self, mock_render, mock_pdf_doc_cls):
+        """When PDF is scanned, extract_from_detection (OCR) is used."""
+        pipeline = self._make_pipeline()
+        detection = self._make_detection()
+        image_bytes = self._make_image_bytes()
+
+        mock_pdf_doc = MagicMock()
+        mock_pdf_doc.is_text_pdf.return_value = False
+        mock_pdf_doc_cls.return_value.__enter__ = MagicMock(return_value=mock_pdf_doc)
+        mock_pdf_doc_cls.return_value.__exit__ = MagicMock(return_value=False)
+
+        pipeline.detector.detect.return_value = [detection]
+        pipeline.extractor.extract_from_detection.return_value = (
+            self._make_extracted_field(raw_text='4.50', normalized='4.50', confidence=0.75)
+        )
+
+        mock_render.return_value = iter([(0, image_bytes)])
+        result = pipeline.process_pdf('/fake/invoice.pdf')
+
+        pipeline.extractor.extract_from_detection.assert_called_once()
+        pipeline.extractor.extract_from_detection_with_pdf.assert_not_called()
+
+    @patch('shared.pdf.extractor.PDFDocument')
+    @patch('shared.pdf.renderer.render_pdf_to_images')
+    def test_pdf_detection_error_falls_back_to_ocr(self, mock_render, mock_pdf_doc_cls):
+        """When PDF text detection throws, fall back to OCR."""
+        pipeline = self._make_pipeline()
+        detection = self._make_detection()
+        image_bytes = self._make_image_bytes()
+
+        mock_ctx = MagicMock()
+        mock_ctx.__enter__ = MagicMock(side_effect=Exception("corrupt PDF"))
+        mock_ctx.__exit__ = MagicMock(return_value=False)
+        mock_pdf_doc_cls.return_value = mock_ctx
+
+        pipeline.detector.detect.return_value = [detection]
+        pipeline.extractor.extract_from_detection.return_value = (
+            self._make_extracted_field(raw_text='4.50', normalized='4.50', confidence=0.75)
+        )
+
+        mock_render.return_value = iter([(0, image_bytes)])
+        result = pipeline.process_pdf('/fake/invoice.pdf')
+
+        pipeline.extractor.extract_from_detection.assert_called_once()
+        pipeline.extractor.extract_from_detection_with_pdf.assert_not_called()
+
+    @patch('shared.pdf.extractor.PDFDocument')
+    @patch('shared.pdf.renderer.render_pdf_to_images')
+    def test_text_pdf_passes_correct_args(self, mock_render, mock_pdf_doc_cls):
+        """Verify correct token list and image dimensions are passed."""
+        from shared.pdf.extractor import Token
+
+        pipeline = self._make_pipeline()
+        detection = self._make_detection()
+        image_bytes = self._make_image_bytes()  # 100x100 PNG
+
+        mock_pdf_doc = MagicMock()
+        mock_pdf_doc.is_text_pdf.return_value = True
+        mock_pdf_doc.page_count = 1
+        tokens = [
+            Token(text="Fakturabelopp:", bbox=(50, 190, 100, 210), page_no=0),
+            Token(text="2.254,50", bbox=(105, 190, 180, 210), page_no=0),
+            Token(text="SEK", bbox=(185, 190, 210, 210), page_no=0),
+        ]
+        mock_pdf_doc.extract_text_tokens.return_value = iter(tokens)
+        mock_pdf_doc_cls.return_value.__enter__ = MagicMock(return_value=mock_pdf_doc)
+        mock_pdf_doc_cls.return_value.__exit__ = MagicMock(return_value=False)
+
+        pipeline.detector.detect.return_value = [detection]
+        pipeline.extractor.extract_from_detection_with_pdf.return_value = (
+            self._make_extracted_field()
+        )
+
+        mock_render.return_value = iter([(0, image_bytes)])
+        pipeline.process_pdf('/fake/invoice.pdf')
+
+        call_args = pipeline.extractor.extract_from_detection_with_pdf.call_args[0]
+        assert call_args[0] == detection
+        assert len(call_args[1]) == 3  # 3 tokens passed
+        assert call_args[2] == 100  # image width
+        assert call_args[3] == 100  # image height
+
+
 if __name__ == '__main__':
     pytest.main([__file__, '-v'])