unstructured-ingest 0.3.13__py3-none-any.whl

test/integration/connectors/test_neo4j.py

@@ -0,0 +1,236 @@
+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
+from test.integration.connectors.utils.docker import container_context
+from unstructured_ingest.error import DestinationConnectionError
+from unstructured_ingest.utils.chunking import elements_from_base64_gzipped_json
+from unstructured_ingest.v2.interfaces.file_data import (
+    FileData,
+    FileDataSourceMetadata,
+    SourceIdentifiers,
+)
+from unstructured_ingest.v2.processes.connectors.neo4j import (
+    CONNECTOR_TYPE,
+    Label,
+    Neo4jAccessConfig,
+    Neo4jConnectionConfig,
+    Neo4jUploader,
+    Neo4jUploaderConfig,
+    Neo4jUploadStager,
+    Relationship,
+)
+
+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)
+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)
+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}) RETURN n"))[0]
+        )
+        document_nodes_count = len(
+            (await driver.execute_query(f"MATCH (n: {Label.DOCUMENT}) RETURN n"))[0]
+        )
+        element_nodes_count = len(
+            (await driver.execute_query(f"MATCH (n: {Label.UNSTRUCTURED_ELEMENT}) 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}]->(:{Label.DOCUMENT}) RETURN r"
+        )
+        part_of_document_count = len(records)
+
+        records, _, _ = await driver.execute_query(
+            f"MATCH (:{Label.CHUNK})-[r:{Relationship.NEXT_CHUNK}]->(:{Label.CHUNK}) 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

@@ -0,0 +1,145 @@
+import os
+
+from test.integration.connectors.utils.validation.source import (
+    SourceValidationConfigs,
+    get_all_file_data,
+    run_all_validations,
+    update_fixtures,
+)
+from unstructured_ingest.v2.interfaces import Downloader, Indexer
+from unstructured_ingest.v2.processes.connectors.notion.connector import (
+    NotionAccessConfig,
+    NotionConnectionConfig,
+    NotionDownloader,
+    NotionDownloaderConfig,
+    NotionIndexer,
+    NotionIndexerConfig,
+)
+
+
+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"],
+        ),
+    )
+
+
+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"],
+        ),
+    )
+
+
+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

@@ -0,0 +1,118 @@
+import os
+import uuid
+from pathlib import Path
+
+import pytest
+from office365.graph_client import GraphClient
+
+from test.integration.connectors.utils.constants import (
+    DESTINATION_TAG,
+)
+from test.integration.utils import requires_env
+from unstructured_ingest.v2.interfaces import FileData, SourceIdentifiers
+from unstructured_ingest.v2.processes.connectors.onedrive import (
+    CONNECTOR_TYPE,
+    OnedriveAccessConfig,
+    OnedriveConnectionConfig,
+    OnedriveUploader,
+    OnedriveUploaderConfig,
+)
+
+
+@pytest.fixture
+@pytest.mark.xfail(
+    reason="Issues with test setup on the provider side."
+)  # TODO: remove line when issues are addressed
+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.tags(CONNECTOR_TYPE, DESTINATION_TAG)
+@requires_env("MS_CLIENT_CRED", "MS_CLIENT_ID", "MS_TENANT_ID", "MS_USER_PNAME")
+@pytest.mark.xfail(
+    reason="Issues with test setup on the provider side."
+)  # TODO: remove line when issues are addressed
+def test_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
+
+    uploaded_file = (
+        drive.root.get_by_path(destination_fullpath).select(["id", "name"]).get().execute_query()
+    )
+
+    # Check if the file exists
+    assert uploaded_file is not None
+    assert uploaded_file.name == upload_file.name