init conftest.py and move fixtures, move csv test to test_split_pdf_hook

Author: Coniferish

_test_unstructured_client/conftest.py

@@ -0,0 +1,19 @@
+from __future__ import annotations
+
+import os
+from pathlib import Path
+from typing import Generator
+import pytest
+
+from unstructured_client.sdk import UnstructuredClient
+
+
+@pytest.fixture(scope="module")
+def client() -> Generator[UnstructuredClient, None, None]:
+    _client = UnstructuredClient(api_key_auth=os.getenv("UNSTRUCTURED_API_KEY"), server='free-api')
+    yield _client
+
+
+@pytest.fixture(scope="module")
+def doc_path() -> Path:
+    return Path(__file__).resolve().parents[1] / "_sample_docs"

_test_unstructured_client/integration/test_integration_freemium.py

@@ -3,29 +3,16 @@
 import asyncio
 import json
 import os
-from pathlib import Path
 
 import pytest
 from deepdiff import DeepDiff
 
 from unstructured_client import UnstructuredClient
 from unstructured_client.models import shared, operations
 from unstructured_client.models.errors import SDKError, ServerError, HTTPValidationError
-from unstructured_client.models.shared.partition_parameters import OutputFormat
 from unstructured_client.utils.retries import BackoffStrategy, RetryConfig
 
 
-@pytest.fixture(scope="module")
-def client() -> UnstructuredClient:
-    _client = UnstructuredClient(api_key_auth=os.getenv("UNSTRUCTURED_API_KEY"), server='free-api')
-    yield _client
-
-
-@pytest.fixture(scope="module")
-def doc_path() -> Path:
-    return Path(__file__).resolve().parents[2] / "_sample_docs"
-
-
 @pytest.mark.parametrize("split_pdf", [True, False])
 @pytest.mark.parametrize("strategy", ["fast", "ocr_only", "hi_res"])
 def test_partition_strategies(split_pdf, strategy, client, doc_path):
@@ -223,29 +210,3 @@ async def call_api():
     uvloop.install()
     elements = asyncio.run(call_api())
     assert len(elements) > 0
-
-
-def test_partition_csv_response(client, doc_path):
-    filename = "layout-parser-paper-fast.pdf"
-    with open(doc_path / filename, "rb") as f:
-        files = shared.Files(
-            content=f.read(),
-            file_name=filename,
-        )
-
-    req = operations.PartitionRequest(
-        partition_parameters=shared.PartitionParameters(
-            files=files,
-            output_format=OutputFormat.TEXT_CSV,
-        )
-    )
-
-    response = client.general.partition(request=req)
-    assert response.status_code == 200
-    assert response.content_type == "text/csv; charset=utf-8"
-    assert response.elements is None
-    assert response.csv_elements is not None
-    assert response.csv_elements.startswith(
-        "type,element_id,text,filetype,languages,page_number,filename,parent_id"
-        "\nTitle,6aa0ff22f91bbe7e26e8e25ca8052acd,Layout"
-    )

_test_unstructured_client/unit/test_request_utils.py

@@ -1,9 +1,10 @@
+from __future__ import annotations
+
 # Get unit tests for request_utils.py module
 import httpx
 import pytest
 
-from unstructured_client._hooks.custom.request_utils import create_pdf_chunk_request_params, get_multipart_stream_fields
-from unstructured_client.models import shared
+from unstructured_client._hooks.custom.request_utils import create_pdf_chunk_request_params, create_response, get_multipart_stream_fields
 
 
 # make the above test using @pytest.mark.parametrize
@@ -30,6 +31,7 @@ def test_get_multipart_stream_fields(input_request, expected):
     fields = get_multipart_stream_fields(input_request)
     assert fields == expected
 
+
 def test_multipart_stream_fields_raises_value_error_when_filename_is_not_set():
     with pytest.raises(ValueError):
         get_multipart_stream_fields(httpx.Request(
@@ -40,6 +42,7 @@ def test_multipart_stream_fields_raises_value_error_when_filename_is_not_set():
             headers={"Content-Type": "multipart/form-data; boundary=----WebKitFormBoundary7MA4YWxkTrZu0gW"}),
         )
 
+
 @pytest.mark.parametrize(("input_form_data", "page_number", "expected_form_data"), [
     (
             {"hello": "world"},
@@ -70,3 +73,24 @@ def test_multipart_stream_fields_raises_value_error_when_filename_is_not_set():
 def test_create_pdf_chunk_request_params(input_form_data, page_number, expected_form_data):
     form_data = create_pdf_chunk_request_params(input_form_data, page_number)
     assert form_data == expected_form_data
+
+
+def test_create_response_for_json():
+    elements = [
+        {"type": "Title", "text": "Hello, World!"},
+        {"type": "NarrativeText", "text": "Goodbye!"},
+    ]
+    response = create_response(elements)
+    assert response.status_code == 200
+    assert response.json() == elements
+    assert response.headers["Content-Type"] == "application/json"
+
+
+def test_create_response_for_csv():
+    elements = 'type,element_id,text,languages,page_number,filename,filetype,parent_id' \
+        '\nTitle,f73329878fbbb0bb131a83e7b6daacbe,Module One - Introduction to Product' \
+        ' Development and Quality Assurance,[\'eng\'],1,list-item-example-1.pdf,application/pdf,'
+    response = create_response(elements)
+    assert response.status_code == 200
+    assert response.json() == None
+    assert response.headers["Content-Type"] == "text/csv"

_test_unstructured_client/unit/test_split_pdf_hook.py

@@ -1,17 +1,13 @@
 from __future__ import annotations
 
 import asyncio
-import io
-import logging
 from asyncio import Task
 from collections import Counter
 from functools import partial
-from typing import Coroutine
 
-import httpx
 import pytest
 import requests
-from requests_toolbelt import MultipartDecoder, MultipartEncoder
+from requests_toolbelt import MultipartDecoder
 
 from unstructured_client._hooks.custom import form_utils, pdf_utils, request_utils
 from unstructured_client._hooks.custom.form_utils import (
@@ -28,7 +24,9 @@
     SplitPdfHook,
     get_optimal_split_size, run_tasks,
 )
-from unstructured_client.models import shared
+from unstructured_client.models import operations, shared
+from unstructured_client.models.shared.partition_parameters import OutputFormat
+from unstructured_client.sdk import UnstructuredClient
 
 
 def test_unit_clear_operation():
@@ -434,3 +432,37 @@ async def test_remaining_tasks_cancelled_when_fails_disallowed():
     await asyncio.sleep(1)
     print("Cancelled amount: ", cancelled_counter["cancelled"])
     assert len(tasks) > cancelled_counter["cancelled"] > 0
+
+
+@pytest.mark.parametrize("split_pdf_page", [True, False])
+def test_integration_get_split_csv_response(split_pdf_page, doc_path):
+    try:
+        response = requests.get("http://127.0.0.1:8000/general/docs")
+        assert response.status_code == 200
+    except requests.exceptions.ConnectionError:
+        assert False, "The unstructured-api is not running on localhost:8000"
+        
+    client = UnstructuredClient(api_key_auth="", server_url="127.0.0.1:8000")
+    filename = "layout-parser-paper.pdf"
+    with open(doc_path / filename, "rb") as f:
+        files = shared.Files(
+            content=f.read(),
+            file_name=filename,
+        )
+    req = operations.PartitionRequest(
+        partition_parameters=shared.PartitionParameters(
+            files=files,
+            output_format=OutputFormat.TEXT_CSV,
+            split_pdf_page=split_pdf_page,
+        )
+    )
+    
+    resp = client.general.partition(request=req)
+
+    assert resp.status_code == 200
+    assert resp.content_type == "text/csv; charset=utf-8"
+    assert resp.elements is None
+    assert resp.csv_elements is not None
+    assert resp.csv_elements.startswith(
+        "type,element_id,text,filetype,languages,page_number,filename,parent_id"
+    )

src/unstructured_client/_hooks/custom/request_utils.py

@@ -207,7 +207,8 @@ def prepare_request_headers(
     new_headers.pop("Content-Length", None)
     return new_headers
 
-def create_response(elements: list) -> httpx.Response:
+
+def create_response(elements: list[dict[str, Any] | bytes]) -> httpx.Response:
     """
     Creates a modified response object with updated content.
 
@@ -218,9 +219,9 @@ def create_response(elements: list) -> httpx.Response:
     Returns:
         The modified response object with updated content.
     """
-    if not isinstance(elements[0], dict):
-        response = httpx.Response(status_code=200, headers={"Content-Type": "text/csv"})
-        content = b''.join(elements)
+    if isinstance(elements, list) and all(isinstance(element, bytes) for element in elements):
+        response = httpx.Response(status_code=200, headers={"Content-Type": "text/csv; charset=utf-8"})
+        content = b''.join(elements) # type: ignore
     else:
         response = httpx.Response(status_code=200, headers={"Content-Type": "application/json"})
         content = json.dumps(elements).encode()