test: before_request split pdf hook unit tests

Author: mpolomdeepsense

_test_unstructured_client/integration/test_integration.py

@@ -7,7 +7,6 @@
 
 import pytest
 from deepdiff import DeepDiff
-from httpx import RequestError
 from unstructured_client import UnstructuredClient
 from unstructured_client.models import shared, operations
 from unstructured_client.models.errors import SDKError, ServerError, HTTPValidationError
@@ -351,52 +350,6 @@ def test_partition_strategy_vlm_anthropic(split_pdf, vlm_model, vlm_model_provid
     assert response.elements[0]["metadata"]["partitioner_type"] == "vlm_partition"
 
 
-@pytest.mark.parametrize(
-    ("pdf_name", "expected_error_message"),
-    [
-        (
-            "failing-encrypted.pdf",
-            "File is encrypted. Please decrypt it with password.",
-        ),
-        (
-            "failing-missing-root.pdf",
-            "File does not appear to be a valid PDF. Error: Cannot find Root object in pdf",
-        ),
-        (
-            "failing-missing-pages.pdf",
-            "File does not appear to be a valid PDF. Error: Invalid object in /Pages",
-        ),
-    ],
-)
-def test_returns_request_error_for_invalid_pdf(
-    caplog: pytest.LogCaptureFixture,
-    doc_path: Path,
-    client: UnstructuredClient,
-    pdf_name: str,
-    expected_error_message: str,
-):
-    """Test that we get a RequestError with the correct error message for invalid PDF files."""
-    with open(doc_path / pdf_name, "rb") as f:
-        files = shared.Files(
-            content=f.read(),
-            file_name=pdf_name,
-        )
-
-    req = operations.PartitionRequest(
-        partition_parameters=shared.PartitionParameters(
-            files=files,
-            strategy="fast",
-            split_pdf_page=True,
-        )
-    )
-
-    with pytest.raises(RequestError) as exc_info:
-        client.general.partition(request=req)
-
-    assert exc_info.value.request is not None
-    assert expected_error_message in caplog.text
-
-
 def test_returns_422_for_invalid_pdf(
     caplog: pytest.LogCaptureFixture,
     doc_path: Path,

_test_unstructured_client/unit/test_split_pdf_hook.py

@@ -8,10 +8,12 @@
 from unittest.mock import MagicMock, patch
 
 import httpx
+from httpx import RequestError
 import pytest
 import requests
 from requests_toolbelt import MultipartDecoder
 
+from _test_unstructured_client.unit_utils import sample_docs_path
 from unstructured_client._hooks.custom import form_utils, pdf_utils, request_utils
 from unstructured_client._hooks.custom.form_utils import (
     FormData,
@@ -29,6 +31,7 @@
     SplitPdfHook,
     get_optimal_split_size, run_tasks,
 )
+from unstructured_client._hooks.types import BeforeRequestContext
 from unstructured_client.models import shared
 
 
@@ -462,3 +465,64 @@ def test_unit_get_split_pdf_cache_tmp_data_dir_uses_dir_from_form_data(mock_path
     mock_path.assert_called_once_with(mock_dir)
     mock_path_instance.exists.assert_called_once()
     assert result == str(Path(mock_dir).resolve())
+
+
+def test_before_request_raises_request_error_when_pdf_check_fails():
+    """Test that before_request raises RequestError when pdf_utils.check_pdf throws PDFValidationError."""
+    hook = SplitPdfHook()
+    
+    # Initialize the hook with a mock client
+    mock_client = MagicMock()
+    hook.sdk_init(base_url="http://localhost:8888", client=mock_client)
+    
+    # Create a mock request context
+    mock_hook_ctx = MagicMock()
+    mock_hook_ctx.operation_id = "partition"
+    
+    # Create a mock request with proper headers and content
+    mock_request = MagicMock()
+    mock_request.headers = {"Content-Type": "multipart/form-data"}
+    mock_request.url.host = "localhost"
+    
+    # Mock the form data to include the necessary fields for PDF splitting
+    mock_pdf_file = MagicMock()
+    mock_pdf_file.read.return_value = b"mock_pdf_content"
+    
+    mock_form_data = {
+        "split_pdf_page": "true",
+        "files": {
+            "filename": "test.pdf",
+            "content_type": "application/pdf",
+            "file": mock_pdf_file
+        }
+    }
+    
+    # Mock the PDF reader object
+    mock_pdf_reader = MagicMock()
+    
+    # Define the error message that will be raised
+    error_message = "File does not appear to be a valid PDF."
+    
+    with patch("unstructured_client._hooks.custom.request_utils.get_multipart_stream_fields") as mock_get_fields, \
+         patch("unstructured_client._hooks.custom.pdf_utils.read_pdf") as mock_read_pdf, \
+         patch("unstructured_client._hooks.custom.pdf_utils.check_pdf") as mock_check_pdf, \
+         patch("unstructured_client._hooks.custom.request_utils.get_base_url") as mock_get_base_url:
+        
+        # Set up the mocks
+        mock_get_fields.return_value = mock_form_data
+        mock_read_pdf.return_value = mock_pdf_reader
+        mock_check_pdf.side_effect = pdf_utils.PDFValidationError(error_message)
+        mock_get_base_url.return_value = "http://localhost:8888"
+        
+        # Call the method under test and verify it raises RequestError
+        with pytest.raises(RequestError) as exc_info:
+            hook.before_request(mock_hook_ctx, mock_request)
+        
+        # Verify the exception has the correct message and request object
+        assert str(exc_info.value) == error_message
+        assert exc_info.value.request == mock_request
+        
+        # Verify that the mocked functions were called as expected
+        mock_get_fields.assert_called_once_with(mock_request)
+        mock_read_pdf.assert_called_once_with(mock_pdf_file)
+        mock_check_pdf.assert_called_once_with(mock_pdf_reader)