fix failing test cases

Author: Bob Strahan

lib/idp_common_pkg/tests/unit/classification/test_classification_service.py

@@ -270,16 +270,11 @@ def test_classify_page_bedrock_model_error(
         mock_get_text.return_value = "This is an invoice for $100"
         mock_invoke.side_effect = Exception("Model error")
 
-        # Call the method
-        result = service.classify_page_bedrock(
-            page_id="1", text_uri="s3://bucket/text.txt"
-        )
-
-        # Verify results
-        assert result.page_id == "1"
-        assert result.classification.doc_type == "unclassified"
-        assert result.classification.confidence == 0.0
-        assert result.classification.metadata["error"] == "Model error"
+        # Call the method and expect exception to be raised
+        with pytest.raises(Exception, match="Model error"):
+            service.classify_page_bedrock(
+                page_id="1", text_uri="s3://bucket/text.txt"
+            )
 
     @patch("boto3.client")
     def test_classify_page_sagemaker_success(self, mock_boto_client, mock_config):

lib/idp_common_pkg/tests/unit/extraction/test_extraction_service.py

@@ -300,16 +300,47 @@ def test_process_document_section_no_pages(
         assert "Section 2 has no page IDs" in result.errors[0]
 
     @patch("idp_common.s3.get_text_content")
+    @patch("idp_common.image.prepare_image")
+    @patch("idp_common.image.prepare_bedrock_image_attachment")
+    @patch("idp_common.bedrock.invoke_model")
+    @patch("idp_common.s3.write_content")
+    @patch("idp_common.utils.merge_metering_data")
     @patch("idp_common.metrics.put_metric")
     def test_process_document_section_missing_page(
-        self, mock_put_metric, mock_get_text_content, service, sample_document
+        self,
+        mock_put_metric,
+        mock_merge_metering,
+        mock_write_content,
+        mock_invoke_model,
+        mock_prepare_bedrock_image,
+        mock_prepare_image,
+        mock_get_text_content,
+        service,
+        sample_document,
     ):
         """Test processing a document section with a missing page."""
         # Add a non-existent page ID to the section
         sample_document.sections[0].page_ids.append("999")
 
-        # Mock text content for existing pages
+        # Mock responses
         mock_get_text_content.side_effect = ["Page 1 text", "Page 2 text"]
+        mock_prepare_image.return_value = b"fake_image_data"
+        mock_prepare_bedrock_image.return_value = {"type": "image", "source": {}}
+        mock_invoke_model.return_value = {
+            "response": {
+                "output": {
+                    "message": {
+                        "content": [
+                            {
+                                "text": '{"invoice_number": "INV-123"}'
+                            }
+                        ]
+                    }
+                }
+            },
+            "metering": {"input_tokens": 100, "output_tokens": 50},
+        }
+        mock_merge_metering.return_value = {"total_tokens": 150}
 
         # Process the section
         result = service.process_document_section(sample_document, "1")
@@ -326,12 +357,9 @@ def test_process_document_section_exception(
         # Mock an exception
         mock_get_text_content.side_effect = Exception("Test exception")
 
-        # Process the section
-        result = service.process_document_section(sample_document, "1")
-
-        # Verify error was added
-        assert len(result.errors) == 1
-        assert "Error processing section 1: Test exception" in result.errors[0]
+        # Process the section and expect exception to be raised
+        with pytest.raises(Exception, match="Test exception"):
+            service.process_document_section(sample_document, "1")
 
     def test_extract_json_code_block(self, service):
         """Test extracting JSON from code block."""

lib/idp_common_pkg/tests/unit/summarization/test_summarization_service.py

@@ -325,11 +325,9 @@ def test_process_text_error(self, mock_invoke_model, service):
         # Configure mock to raise exception
         mock_invoke_model.side_effect = Exception("Test error")
 
-        result = service.process_text("Test document text")
-
-        # Verify result
-        assert result.content["error"] == "Error generating summary"
-        assert result.metadata["error"] == "Test error"
+        # Expect exception to be raised
+        with pytest.raises(Exception, match="Test error"):
+            service.process_text("Test document text")
 
     @patch("idp_common.s3.get_text_content")
     @patch("idp_common.s3.write_content")