unstructured-ingest 0.7.1__py3-none-any.whl → 1.0.1__py3-none-any.whl

test/integration/connectors/test_neo4j.py

@@ -1,244 +0,0 @@
-import json
-import time
-import uuid
-from datetime import datetime
-from pathlib import Path
-
-import pytest
-from neo4j import AsyncGraphDatabase, Driver, GraphDatabase
-from neo4j.exceptions import ServiceUnavailable
-from pytest_check import check
-
-from test.integration.connectors.utils.constants import DESTINATION_TAG, GRAPH_DB_TAG
-from test.integration.connectors.utils.docker import container_context
-from unstructured_ingest.data_types.file_data import (
-    FileData,
-    FileDataSourceMetadata,
-    SourceIdentifiers,
-)
-from unstructured_ingest.error import DestinationConnectionError
-from unstructured_ingest.processes.connectors.neo4j import (
-    CONNECTOR_TYPE,
-    Label,
-    Neo4jAccessConfig,
-    Neo4jConnectionConfig,
-    Neo4jUploader,
-    Neo4jUploaderConfig,
-    Neo4jUploadStager,
-    Relationship,
-)
-from unstructured_ingest.utils.chunking import elements_from_base64_gzipped_json
-
-USERNAME = "neo4j"
-PASSWORD = "password"
-URI = "neo4j://localhost:7687"
-DATABASE = "neo4j"
-
-EXPECTED_DOCUMENT_COUNT = 1
-
-
-# NOTE: Precheck tests are read-only so we utilize the same container for all tests.
-# If new tests require clean neo4j container, this fixture's scope should be adjusted.
-@pytest.fixture(autouse=True, scope="module")
-def _neo4j_server():
-    with container_context(
-        image="neo4j:latest", environment={"NEO4J_AUTH": "neo4j/password"}, ports={"7687": "7687"}
-    ):
-        driver = GraphDatabase.driver(uri=URI, auth=(USERNAME, PASSWORD))
-        wait_for_connection(driver)
-        driver.close()
-        yield
-
-
-@pytest.mark.asyncio
-@pytest.mark.tags(DESTINATION_TAG, CONNECTOR_TYPE, GRAPH_DB_TAG)
-async def test_neo4j_destination(upload_file: Path, tmp_path: Path):
-    stager = Neo4jUploadStager()
-    uploader = Neo4jUploader(
-        connection_config=Neo4jConnectionConfig(
-            access_config=Neo4jAccessConfig(password=PASSWORD),  # type: ignore
-            username=USERNAME,
-            uri=URI,
-            database=DATABASE,
-        ),
-        upload_config=Neo4jUploaderConfig(),
-    )
-    file_data = FileData(
-        identifier="mock-file-data",
-        connector_type="neo4j",
-        source_identifiers=SourceIdentifiers(
-            filename=upload_file.name,
-            fullpath=upload_file.name,
-        ),
-        metadata=FileDataSourceMetadata(
-            date_created=str(datetime(2022, 1, 1).timestamp()),
-            date_modified=str(datetime(2022, 1, 2).timestamp()),
-        ),
-    )
-    staged_filepath = stager.run(
-        upload_file,
-        file_data=file_data,
-        output_dir=tmp_path,
-        output_filename=upload_file.name,
-    )
-
-    await uploader.run_async(staged_filepath, file_data)
-    await validate_uploaded_graph(upload_file)
-
-    modified_upload_file = tmp_path / f"modified-{upload_file.name}"
-    with open(upload_file) as file:
-        elements = json.load(file)
-        for element in elements:
-            element["element_id"] = str(uuid.uuid4())
-
-    with open(modified_upload_file, "w") as file:
-        json.dump(elements, file, indent=4)
-
-    staged_filepath = stager.run(
-        modified_upload_file,
-        file_data=file_data,
-        output_dir=tmp_path,
-        output_filename=modified_upload_file.name,
-    )
-    await uploader.run_async(staged_filepath, file_data)
-    await validate_uploaded_graph(modified_upload_file)
-
-
-@pytest.mark.tags(DESTINATION_TAG, CONNECTOR_TYPE, GRAPH_DB_TAG)
-class TestPrecheck:
-    @pytest.fixture
-    def configured_uploader(self) -> Neo4jUploader:
-        return Neo4jUploader(
-            connection_config=Neo4jConnectionConfig(
-                access_config=Neo4jAccessConfig(password=PASSWORD),  # type: ignore
-                username=USERNAME,
-                uri=URI,
-                database=DATABASE,
-            ),
-            upload_config=Neo4jUploaderConfig(),
-        )
-
-    def test_succeeds(self, configured_uploader: Neo4jUploader):
-        configured_uploader.precheck()
-
-    def test_fails_on_invalid_password(self, configured_uploader: Neo4jUploader):
-        configured_uploader.connection_config.access_config.get_secret_value().password = (
-            "invalid-password"
-        )
-        with pytest.raises(
-            DestinationConnectionError,
-            match="{code: Neo.ClientError.Security.Unauthorized}",
-        ):
-            configured_uploader.precheck()
-
-    def test_fails_on_invalid_username(self, configured_uploader: Neo4jUploader):
-        configured_uploader.connection_config.username = "invalid-username"
-        with pytest.raises(
-            DestinationConnectionError, match="{code: Neo.ClientError.Security.Unauthorized}"
-        ):
-            configured_uploader.precheck()
-
-    @pytest.mark.parametrize(
-        ("uri", "expected_error_msg"),
-        [
-            ("neo4j://localhst:7687", "Cannot resolve address"),
-            ("neo4j://localhost:7777", "Unable to retrieve routing information"),
-        ],
-    )
-    def test_fails_on_invalid_uri(
-        self, configured_uploader: Neo4jUploader, uri: str, expected_error_msg: str
-    ):
-        configured_uploader.connection_config.uri = uri
-        with pytest.raises(DestinationConnectionError, match=expected_error_msg):
-            configured_uploader.precheck()
-
-    def test_fails_on_invalid_database(self, configured_uploader: Neo4jUploader):
-        configured_uploader.connection_config.database = "invalid-database"
-        with pytest.raises(
-            DestinationConnectionError, match="{code: Neo.ClientError.Database.DatabaseNotFound}"
-        ):
-            configured_uploader.precheck()
-
-
-def wait_for_connection(driver: Driver, retries: int = 10, delay_seconds: int = 2):
-    attempts = 0
-    while attempts < retries:
-        try:
-            driver.verify_connectivity()
-            return
-        except ServiceUnavailable:
-            time.sleep(delay_seconds)
-            attempts += 1
-
-    pytest.fail("Failed to connect with Neo4j server.")
-
-
-async def validate_uploaded_graph(upload_file: Path):
-    with open(upload_file) as file:
-        elements = json.load(file)
-
-    for element in elements:
-        if "orig_elements" in element["metadata"]:
-            element["metadata"]["orig_elements"] = elements_from_base64_gzipped_json(
-                element["metadata"]["orig_elements"]
-            )
-        else:
-            element["metadata"]["orig_elements"] = []
-
-    expected_chunks_count = len(elements)
-    expected_element_count = len(
-        {
-            origin_element["element_id"]
-            for chunk in elements
-            for origin_element in chunk["metadata"]["orig_elements"]
-        }
-    )
-    expected_nodes_count = expected_chunks_count + expected_element_count + EXPECTED_DOCUMENT_COUNT
-
-    driver = AsyncGraphDatabase.driver(uri=URI, auth=(USERNAME, PASSWORD))
-    try:
-        nodes_count = len((await driver.execute_query("MATCH (n) RETURN n"))[0])
-        chunk_nodes_count = len(
-            (await driver.execute_query(f"MATCH (n: {Label.CHUNK.value}) RETURN n"))[0]
-        )
-        document_nodes_count = len(
-            (await driver.execute_query(f"MATCH (n: {Label.DOCUMENT.value}) RETURN n"))[0]
-        )
-        element_nodes_count = len(
-            (await driver.execute_query(f"MATCH (n: {Label.UNSTRUCTURED_ELEMENT.value}) RETURN n"))[
-                0
-            ]
-        )
-        with check:
-            assert nodes_count == expected_nodes_count
-        with check:
-            assert document_nodes_count == EXPECTED_DOCUMENT_COUNT
-        with check:
-            assert chunk_nodes_count == expected_chunks_count
-        with check:
-            assert element_nodes_count == expected_element_count
-
-        records, _, _ = await driver.execute_query(
-            f"""
-            MATCH ()-[r:{Relationship.PART_OF_DOCUMENT.value}]->(:{Label.DOCUMENT.value})
-            RETURN r
-            """
-        )
-        part_of_document_count = len(records)
-
-        records, _, _ = await driver.execute_query(
-            f"""
-            MATCH (:{Label.CHUNK.value})-[r:{Relationship.NEXT_CHUNK.value}]->(:{Label.CHUNK.value})
-            RETURN r
-            """
-        )
-        next_chunk_count = len(records)
-
-        if not check.any_failures():
-            with check:
-                assert part_of_document_count == expected_chunks_count + expected_element_count
-            with check:
-                assert next_chunk_count == expected_chunks_count - 1
-
-    finally:
-        await driver.close()

test/integration/connectors/test_notion.py

@@ -1,152 +0,0 @@
-import os
-
-import pytest
-
-from test.integration.connectors.utils.constants import SOURCE_TAG, UNCATEGORIZED_TAG
-from test.integration.connectors.utils.validation.source import (
-    SourceValidationConfigs,
-    get_all_file_data,
-    run_all_validations,
-    update_fixtures,
-)
-from unstructured_ingest.interfaces import Downloader, Indexer
-from unstructured_ingest.processes.connectors.notion.connector import (
-    CONNECTOR_TYPE,
-    NotionAccessConfig,
-    NotionConnectionConfig,
-    NotionDownloader,
-    NotionDownloaderConfig,
-    NotionIndexer,
-    NotionIndexerConfig,
-)
-
-
-@pytest.mark.tags(SOURCE_TAG, CONNECTOR_TYPE, UNCATEGORIZED_TAG)
-def test_notion_source_database(temp_dir):
-    # Retrieve environment variables
-    notion_api_key = os.environ["NOTION_API_KEY"]
-
-    # Create connection and indexer configurations
-    access_config = NotionAccessConfig(notion_api_key=notion_api_key)
-    connection_config = NotionConnectionConfig(
-        access_config=access_config,
-    )
-    index_config = NotionIndexerConfig(
-        database_ids=["1722c3765a0a8082b382ebc2c62d3f4c"], recursive=False
-    )
-
-    download_config = NotionDownloaderConfig(download_dir=temp_dir)
-
-    # Instantiate indexer and downloader
-    indexer = NotionIndexer(
-        connection_config=connection_config,
-        index_config=index_config,
-    )
-    downloader = NotionDownloader(
-        connection_config=connection_config,
-        download_config=download_config,
-    )
-
-    # Run the source connector validation
-    source_connector_validation(
-        indexer=indexer,
-        downloader=downloader,
-        configs=SourceValidationConfigs(
-            test_id="notion_database",
-            expected_num_files=1,
-            validate_downloaded_files=True,
-            exclude_fields_extend=["metadata.date_created", "metadata.date_modified"],
-        ),
-    )
-
-
-@pytest.mark.tags(SOURCE_TAG, CONNECTOR_TYPE, UNCATEGORIZED_TAG)
-def test_notion_source_page(temp_dir):
-    # Retrieve environment variables
-    notion_api_key = os.environ["NOTION_API_KEY"]
-
-    # Create connection and indexer configurations
-    access_config = NotionAccessConfig(notion_api_key=notion_api_key)
-    connection_config = NotionConnectionConfig(
-        access_config=access_config,
-    )
-    index_config = NotionIndexerConfig(
-        page_ids=["1572c3765a0a806299f0dd6999f9e4c7"], recursive=False
-    )
-
-    download_config = NotionDownloaderConfig(download_dir=temp_dir)
-
-    # Instantiate indexer and downloader
-    indexer = NotionIndexer(
-        connection_config=connection_config,
-        index_config=index_config,
-    )
-    downloader = NotionDownloader(
-        connection_config=connection_config,
-        download_config=download_config,
-    )
-
-    # Run the source connector validation
-    source_connector_validation(
-        indexer=indexer,
-        downloader=downloader,
-        configs=SourceValidationConfigs(
-            test_id="notion_page",
-            expected_num_files=1,
-            validate_downloaded_files=True,
-            exclude_fields_extend=["metadata.date_created", "metadata.date_modified"],
-        ),
-    )
-
-
-@pytest.mark.tags(SOURCE_TAG, CONNECTOR_TYPE, UNCATEGORIZED_TAG)
-def source_connector_validation(
-    indexer: Indexer,
-    downloader: Downloader,
-    configs: SourceValidationConfigs,
-    overwrite_fixtures: bool = os.getenv("OVERWRITE_FIXTURES", "False").lower() == "true",
-) -> None:
-    # Run common validations on the process of running a source connector, supporting dynamic
-    # validators that get passed in along with comparisons on the saved expected values.
-    # If overwrite_fixtures is st to True, will ignore all validators but instead overwrite the
-    # expected values with what gets generated by this test.
-    all_predownload_file_data = []
-    all_postdownload_file_data = []
-    indexer.precheck()
-    download_dir = downloader.download_config.download_dir
-    test_output_dir = configs.test_output_dir()
-
-    for file_data in indexer.run():
-        assert file_data
-        predownload_file_data = file_data.model_copy(deep=True)
-        all_predownload_file_data.append(predownload_file_data)
-        resp = downloader.run(file_data=file_data)
-        if isinstance(resp, list):
-            for r in resp:
-                postdownload_file_data = r["file_data"].model_copy(deep=True)
-                all_postdownload_file_data.append(postdownload_file_data)
-        else:
-            postdownload_file_data = resp["file_data"].model_copy(deep=True)
-            all_postdownload_file_data.append(postdownload_file_data)
-
-    if not overwrite_fixtures:
-        print("Running validation")
-        run_all_validations(
-            configs=configs,
-            predownload_file_data=all_predownload_file_data,
-            postdownload_file_data=all_postdownload_file_data,
-            download_dir=download_dir,
-            test_output_dir=test_output_dir,
-        )
-    else:
-        print("Running fixtures update")
-        update_fixtures(
-            output_dir=test_output_dir,
-            download_dir=download_dir,
-            all_file_data=get_all_file_data(
-                all_predownload_file_data=all_predownload_file_data,
-                all_postdownload_file_data=all_postdownload_file_data,
-            ),
-            save_downloads=configs.validate_downloaded_files,
-            save_filedata=configs.validate_file_data,
-        )

test/integration/connectors/test_onedrive.py

@@ -1,163 +0,0 @@
-import os
-import uuid
-from pathlib import Path
-
-import pytest
-from office365.graph_client import GraphClient
-
-from test.integration.connectors.utils.constants import (
-    BLOB_STORAGE_TAG,
-    DESTINATION_TAG,
-    SOURCE_TAG,
-)
-from test.integration.connectors.utils.validation.source import (
-    SourceValidationConfigs,
-    source_connector_validation,
-)
-from test.integration.utils import requires_env
-from unstructured_ingest.data_types.file_data import FileData, SourceIdentifiers
-from unstructured_ingest.processes.connectors.onedrive import (
-    CONNECTOR_TYPE,
-    OnedriveAccessConfig,
-    OnedriveConnectionConfig,
-    OnedriveDownloader,
-    OnedriveDownloaderConfig,
-    OnedriveIndexer,
-    OnedriveIndexerConfig,
-    OnedriveUploader,
-    OnedriveUploaderConfig,
-)
-
-
-@pytest.fixture
-def onedrive_test_folder() -> str:
-    """
-    Pytest fixture that creates a test folder in OneDrive and deletes it after test run.
-    """
-    connection_config = get_connection_config()
-    user_pname = connection_config.user_pname
-
-    # Get the OneDrive client
-    client: GraphClient = connection_config.get_client()
-    drive = client.users[user_pname].drive
-
-    # Generate a unique test folder path
-    test_folder_path = f"utic-test-output-{uuid.uuid4()}"
-
-    # Create the test folder
-    root = drive.root
-    folder = root.create_folder(test_folder_path).execute_query()
-    print(f"created folder: {folder.name}")
-    try:
-        yield test_folder_path
-    finally:
-        # Teardown: delete the test folder and its contents
-        folder.delete_object().execute_query()
-        print(f"successfully deleted folder: {folder.name}")
-
-
-def get_connection_config():
-    """
-    Pytest fixture that provides the OnedriveConnectionConfig for tests.
-    """
-    client_id = os.getenv("MS_CLIENT_ID")
-    client_secret = os.getenv("MS_CLIENT_CRED")
-    tenant_id = os.getenv("MS_TENANT_ID")
-    user_pname = os.getenv("MS_USER_PNAME")
-
-    connection_config = OnedriveConnectionConfig(
-        client_id=client_id,
-        tenant=tenant_id,
-        user_pname=user_pname,
-        access_config=OnedriveAccessConfig(client_cred=client_secret),
-    )
-    return connection_config
-
-
-@pytest.mark.asyncio
-@pytest.mark.tags(CONNECTOR_TYPE, SOURCE_TAG, BLOB_STORAGE_TAG)
-@requires_env("MS_CLIENT_CRED", "MS_CLIENT_ID", "MS_TENANT_ID", "MS_USER_PNAME")
-async def test_onedrive_source(temp_dir):
-    connection_config = get_connection_config()
-    index_config = OnedriveIndexerConfig(recursive=True, path="eml")
-
-    download_config = OnedriveDownloaderConfig(download_dir=temp_dir)
-
-    # Instantiate indexer and downloader
-    indexer = OnedriveIndexer(
-        connection_config=connection_config,
-        index_config=index_config,
-    )
-    downloader = OnedriveDownloader(
-        connection_config=connection_config,
-        download_config=download_config,
-    )
-
-    # Run the source connector validation
-    await source_connector_validation(
-        indexer=indexer,
-        downloader=downloader,
-        configs=SourceValidationConfigs(
-            test_id="onedrive",
-            expected_num_files=1,
-            validate_downloaded_files=True,
-            exclude_fields_extend=[
-                "metadata.date_created",
-                "metadata.date_modified",
-                "additional_metadata.LastModified",
-                "additional_metadata.@microsoft.graph.downloadUrl",
-            ],
-        ),
-    )
-
-
-@pytest.mark.tags(CONNECTOR_TYPE, DESTINATION_TAG, BLOB_STORAGE_TAG)
-@requires_env("MS_CLIENT_CRED", "MS_CLIENT_ID", "MS_TENANT_ID", "MS_USER_PNAME")
-def xtest_onedrive_destination(upload_file: Path, onedrive_test_folder: str):
-    """
-    Integration test for the OneDrive destination connector.
-
-    This test uploads a file to OneDrive and verifies that it exists.
-    """
-    connection_config = get_connection_config()
-    # Retrieve user principal name from the connection config
-    user_pname = connection_config.user_pname
-
-    # The test folder is provided by the fixture
-    destination_folder = onedrive_test_folder
-    destination_fullpath = f"{destination_folder}/{upload_file.name}"
-
-    # Configure the uploader with remote_url
-    upload_config = OnedriveUploaderConfig(remote_url=f"onedrive://{destination_folder}")
-
-    uploader = OnedriveUploader(
-        connection_config=connection_config,
-        upload_config=upload_config,
-    )
-
-    file_data = FileData(
-        source_identifiers=SourceIdentifiers(
-            fullpath=destination_fullpath,
-            filename=upload_file.name,
-        ),
-        connector_type=CONNECTOR_TYPE,
-        identifier="mock_file_data",
-    )
-    uploader.precheck()
-    uploader.run(path=upload_file, file_data=file_data)
-
-    # Verify that the file was uploaded
-    client = connection_config.get_client()
-    drive = client.users[user_pname].drive
-
-    # Workaround: File should not have .json in the metadata.filename it comes from embedder
-    uploaded_file = (
-        drive.root.get_by_path(f"{destination_fullpath}.json")
-        .select(["id", "name"])
-        .get()
-        .execute_query()
-    )
-
-    # Check if the file exists
-    assert uploaded_file is not None
-    assert uploaded_file.name == f"{upload_file.name}.json"