test: revive ingest unit tests (#688)

Author: ryannikolaidis

.github/workflows/ci.yml

@@ -167,7 +167,11 @@ jobs:
         source .venv/bin/activate
         mkdir "$NLTK_DATA"
         make install-ci
-    - name: Test
+    - name: Test Ingest (unit)
+      run: |
+        source .venv/bin/activate
+        PYTHONPATH=. pytest test_unstructured_ingest/unit
+    - name: Test (end-to-end)
       env:
         GH_READ_ONLY_ACCESS_TOKEN: ${{ secrets.GH_READ_ONLY_ACCESS_TOKEN }}
         SLACK_TOKEN: ${{ secrets.SLACK_TOKEN }}

test_unstructured_ingest/test_interfaces.py

@@ -0,0 +1,139 @@
+import os
+import pathlib
+from dataclasses import dataclass
+
+import pytest
+
+from unstructured.ingest.connector.local import LocalIngestDoc, SimpleLocalConfig
+from unstructured.ingest.interfaces import (
+    BaseConnector,
+    BaseConnectorConfig,
+    BaseIngestDoc,
+    StandardConnectorConfig,
+)
+from unstructured.partition.auto import partition
+from unstructured.staging.base import convert_to_dict
+
+DIRECTORY = pathlib.Path(__file__).parent.resolve()
+EXAMPLE_DOCS_DIRECTORY = os.path.join(DIRECTORY, "../..", "example-docs")
+TEST_DOWNLOAD_DIR="/tmp"
+TEST_OUTPUT_DIR="/tmp"
+TEST_ID="test"
+TEST_FILE_PATH=os.path.join(EXAMPLE_DOCS_DIRECTORY, "book-war-and-peace-1p.txt")
+
+@dataclass
+class TestConfig(BaseConnectorConfig):
+    id: str
+    path: str
+
+TEST_CONFIG=TestConfig(id=TEST_ID, path=TEST_FILE_PATH)
+
+@dataclass
+class TestIngestDoc(BaseIngestDoc):
+    config: TestConfig
+
+    @property
+    def filename(self):
+        return "test"
+    
+    def cleanup_file(self):
+        pass
+
+    def get_file(self):
+        pass
+
+    def has_output(self):
+        return True
+    
+    def write_result(self, result):
+        pass
+
+@pytest.fixture()
+def partition_test_results():
+    # Reusable partition test results, calculated only once
+    return partition(filename=str(TEST_FILE_PATH))
+
+@pytest.fixture()
+def partition_file_test_results(partition_test_results):
+    # Reusable partition_file test results, calculated only once
+    return convert_to_dict(partition_test_results)
+
+def test_process_file_fields_include_default(mocker, partition_test_results):
+    """Validate when metadata_include and metadata_exclude are not set, all fields:
+    ("element_id", "text", "type", "metadata") are included"""
+    mocker.patch(
+        "unstructured.ingest.interfaces.partition",
+        return_value=partition_test_results,
+    )
+    test_ingest_doc = TestIngestDoc(
+        config=TEST_CONFIG,
+        standard_config=StandardConnectorConfig(
+            download_dir=TEST_DOWNLOAD_DIR,
+            output_dir=TEST_OUTPUT_DIR,
+            metadata_include="filename,page_number",
+        ),
+    )
+    isd_elems = test_ingest_doc.process_file()
+    assert len(isd_elems)
+    for elem in isd_elems:
+        assert {"element_id", "text", "type", "metadata"} == set(elem.keys())
+
+
+def test_process_file_metadata_includes_filename_and_page_number(mocker, partition_test_results):
+    """Validate when metadata_include is set to "filename,page_number",
+    only filename is included in metadata"""
+    mocker.patch(
+        "unstructured.ingest.interfaces.partition",
+        return_value=partition_test_results,
+    )
+    test_ingest_doc = TestIngestDoc(
+        config=TEST_CONFIG,
+        standard_config=StandardConnectorConfig(
+            download_dir=TEST_DOWNLOAD_DIR,
+            output_dir=TEST_OUTPUT_DIR,
+            metadata_include="filename,page_number",
+        ),
+    )
+    isd_elems = test_ingest_doc.process_file()
+    assert len(isd_elems)
+    for elem in isd_elems:
+        assert set(elem["metadata"].keys()) == {"filename", "page_number"}
+
+def test_process_file_metadata_exclude_filename_pagenum(mocker, partition_test_results):
+    """Validate when metadata_exclude is set to "filename,page_number",
+    neither filename nor page_number are included in metadata"""
+    mocker.patch(
+        "unstructured.ingest.interfaces.partition",
+        return_value=partition_test_results,
+    )
+    test_ingest_doc = TestIngestDoc(
+        config=TEST_CONFIG,
+        standard_config=StandardConnectorConfig(
+            download_dir=TEST_DOWNLOAD_DIR,
+            output_dir=TEST_OUTPUT_DIR,
+            metadata_exclude="filename,page_number",
+        ),
+    )
+    isd_elems = test_ingest_doc.process_file()
+    assert len(isd_elems)
+    for elem in isd_elems:
+        assert "filename" not in elem["metadata"].keys()
+        assert "page_number" not in elem["metadata"].keys()
+
+def test_process_file_flatten_metadata(mocker, partition_test_results):
+    mocker.patch(
+        "unstructured.ingest.interfaces.partition",
+        return_value=partition_test_results,
+    )
+    test_ingest_doc = TestIngestDoc(
+        config=TEST_CONFIG,
+        standard_config=StandardConnectorConfig(
+            download_dir=TEST_DOWNLOAD_DIR,
+            output_dir=TEST_OUTPUT_DIR,
+            metadata_include="filename,page_number",
+            flatten_metadata=True,
+        ),
+    )
+    isd_elems = test_ingest_doc.process_file()
+    for elem in isd_elems:
+        assert {"element_id", "text", "type", "filename", "page_number"} == set(elem.keys())