unstructured-ingest 0.3.7__py3-none-any.whl → 0.3.9__py3-none-any.whl

test/integration/connectors/sql/test_sqlite.py

@@ -1,14 +1,18 @@
+import json
 import sqlite3
 import tempfile
-from contextlib import contextmanager
 from pathlib import Path
 
-import pandas as pd
 import pytest
+from _pytest.fixtures import TopRequest
 
 from test.integration.connectors.utils.constants import DESTINATION_TAG, SOURCE_TAG, env_setup_path
-from test.integration.connectors.utils.validation import (
-    ValidationConfigs,
+from test.integration.connectors.utils.validation.destination import (
+    StagerValidationConfigs,
+    stager_validation,
+)
+from test.integration.connectors.utils.validation.source import (
+    SourceValidationConfigs,
     source_connector_validation,
 )
 from unstructured_ingest.v2.interfaces import FileData
@@ -26,8 +30,8 @@ from unstructured_ingest.v2.processes.connectors.sql.sqlite import (
 SEED_DATA_ROWS = 20
 
 
-@contextmanager
-def sqlite_download_setup() -> Path:
+@pytest.fixture
+def source_database_setup() -> Path:
     with tempfile.TemporaryDirectory() as tmpdir:
         db_path = Path(tmpdir) / "mock_database.db"
         db_init_path = env_setup_path / "sql" / "sqlite" / "source" / "sqlite-schema.sql"
@@ -49,49 +53,42 @@ def sqlite_download_setup() -> Path:
 
 @pytest.mark.asyncio
 @pytest.mark.tags(CONNECTOR_TYPE, SOURCE_TAG, "sql")
-async def test_sqlite_source():
-    with sqlite_download_setup() as db_path:
-        with tempfile.TemporaryDirectory() as tmpdir:
-            connection_config = SQLiteConnectionConfig(database_path=db_path)
-            indexer = SQLiteIndexer(
-                connection_config=connection_config,
-                index_config=SQLiteIndexerConfig(
-                    table_name="cars", id_column="car_id", batch_size=5
-                ),
-            )
-            downloader = SQLiteDownloader(
-                connection_config=connection_config,
-                download_config=SQLiteDownloaderConfig(
-                    fields=["car_id", "brand"], download_dir=Path(tmpdir)
-                ),
-            )
-            await source_connector_validation(
-                indexer=indexer,
-                downloader=downloader,
-                configs=ValidationConfigs(
-                    test_id="sqlite",
-                    expected_num_files=SEED_DATA_ROWS,
-                    expected_number_indexed_file_data=4,
-                    validate_downloaded_files=True,
-                ),
-            )
-
-
-@contextmanager
-def sqlite_upload_setup() -> Path:
+async def test_sqlite_source(source_database_setup: Path, temp_dir: Path):
+    connection_config = SQLiteConnectionConfig(database_path=source_database_setup)
+    indexer = SQLiteIndexer(
+        connection_config=connection_config,
+        index_config=SQLiteIndexerConfig(table_name="cars", id_column="car_id", batch_size=5),
+    )
+    downloader = SQLiteDownloader(
+        connection_config=connection_config,
+        download_config=SQLiteDownloaderConfig(fields=["car_id", "brand"], download_dir=temp_dir),
+    )
+    await source_connector_validation(
+        indexer=indexer,
+        downloader=downloader,
+        configs=SourceValidationConfigs(
+            test_id="sqlite",
+            expected_num_files=SEED_DATA_ROWS,
+            expected_number_indexed_file_data=4,
+            validate_downloaded_files=True,
+        ),
+    )
+
+
+@pytest.fixture
+def destination_database_setup(temp_dir: Path) -> Path:
     # Provision the local file that sqlite points to to have the desired schema for the integration
     # tests and make sure the file and connection get cleaned up by using a context manager.
-    with tempfile.TemporaryDirectory() as tmpdir:
-        db_path = Path(tmpdir) / "elements.db"
-        db_init_path = env_setup_path / "sql" / "sqlite" / "destination" / "sqlite-schema.sql"
-        assert db_init_path.exists()
-        assert db_init_path.is_file()
-        with sqlite3.connect(database=db_path) as sqlite_connection:
-            with db_init_path.open("r") as f:
-                query = f.read()
-            cursor = sqlite_connection.cursor()
-            cursor.executescript(query)
-        yield db_path
+    db_path = temp_dir / "elements.db"
+    db_init_path = env_setup_path / "sql" / "sqlite" / "destination" / "sqlite-schema.sql"
+    assert db_init_path.exists()
+    assert db_init_path.is_file()
+    with sqlite3.connect(database=db_path) as sqlite_connection:
+        with db_init_path.open("r") as f:
+            query = f.read()
+        cursor = sqlite_connection.cursor()
+        cursor.executescript(query)
+    return db_path
 
 
 def validate_destination(db_path: Path, expected_num_elements: int):
@@ -114,34 +111,48 @@ def validate_destination(db_path: Path, expected_num_elements: int):
 
 @pytest.mark.asyncio
 @pytest.mark.tags(CONNECTOR_TYPE, DESTINATION_TAG, "sql")
-async def test_sqlite_destination(upload_file: Path):
+async def test_sqlite_destination(
+    upload_file: Path, temp_dir: Path, destination_database_setup: Path
+):
     # the sqlite destination connector doesn't leverage the file data but is required as an input,
     # mocking it with arbitrary values to meet the base requirements:
     mock_file_data = FileData(identifier="mock file data", connector_type=CONNECTOR_TYPE)
-    with sqlite_upload_setup() as db_path:
-        with tempfile.TemporaryDirectory() as tmpdir:
-            stager = SQLiteUploadStager()
-            stager_params = {
-                "elements_filepath": upload_file,
-                "file_data": mock_file_data,
-                "output_dir": Path(tmpdir),
-                "output_filename": "test_db",
-            }
-            if stager.is_async():
-                staged_path = await stager.run_async(**stager_params)
-            else:
-                staged_path = stager.run(**stager_params)
-
-            # The stager should append the `.json` suffix to the output filename passed in.
-            assert staged_path.name == "test_db.json"
-
-            uploader = SQLiteUploader(
-                connection_config=SQLiteConnectionConfig(database_path=db_path)
-            )
-            uploader.run(path=staged_path, file_data=mock_file_data)
-
-            staged_df = pd.read_json(staged_path, orient="records", lines=True)
-            validate_destination(db_path=db_path, expected_num_elements=len(staged_df))
-
-            uploader.run(path=staged_path, file_data=mock_file_data)
-            validate_destination(db_path=db_path, expected_num_elements=len(staged_df))
+    stager = SQLiteUploadStager()
+    staged_path = stager.run(
+        elements_filepath=upload_file,
+        file_data=mock_file_data,
+        output_dir=temp_dir,
+        output_filename=upload_file.name,
+    )
+
+    # The stager should append the `.json` suffix to the output filename passed in.
+    assert staged_path.suffix == upload_file.suffix
+
+    uploader = SQLiteUploader(
+        connection_config=SQLiteConnectionConfig(database_path=destination_database_setup)
+    )
+    uploader.precheck()
+    uploader.run(path=staged_path, file_data=mock_file_data)
+
+    with staged_path.open("r") as f:
+        staged_data = json.load(f)
+    validate_destination(db_path=destination_database_setup, expected_num_elements=len(staged_data))
+
+    uploader.run(path=staged_path, file_data=mock_file_data)
+    validate_destination(db_path=destination_database_setup, expected_num_elements=len(staged_data))
+
+
+@pytest.mark.parametrize("upload_file_str", ["upload_file_ndjson", "upload_file"])
+def test_sqlite_stager(
+    request: TopRequest,
+    upload_file_str: str,
+    tmp_path: Path,
+):
+    upload_file: Path = request.getfixturevalue(upload_file_str)
+    stager = SQLiteUploadStager()
+    stager_validation(
+        configs=StagerValidationConfigs(test_id=CONNECTOR_TYPE, expected_count=22),
+        input_file=upload_file,
+        stager=stager,
+        tmp_dir=tmp_path,
+    )

test/integration/connectors/test_astradb.py

@@ -5,10 +5,15 @@ from pathlib import Path
 from uuid import uuid4
 
 import pytest
+from _pytest.fixtures import TopRequest
 from astrapy import Collection
 from astrapy import DataAPIClient as AstraDBClient
 
 from test.integration.connectors.utils.constants import DESTINATION_TAG, SOURCE_TAG
+from test.integration.connectors.utils.validation.destination import (
+    StagerValidationConfigs,
+    stager_validation,
+)
 from test.integration.utils import requires_env
 from unstructured_ingest.v2.interfaces import FileData, SourceIdentifiers
 from unstructured_ingest.v2.processes.connectors.astradb import (
@@ -108,7 +113,7 @@ def collection(upload_file: Path) -> Collection:
 @pytest.mark.asyncio
 @pytest.mark.tags(CONNECTOR_TYPE, DESTINATION_TAG)
 @requires_env("ASTRA_DB_API_ENDPOINT", "ASTRA_DB_APPLICATION_TOKEN")
-async def test_azure_ai_search_destination(
+async def test_astra_search_destination(
     upload_file: Path,
     collection: Collection,
     tmp_path: Path,
@@ -154,3 +159,19 @@ async def test_azure_ai_search_destination(
         f"Expected count ({expected_count}) doesn't match how "
         f"much came back from collection: {current_count}"
     )
+
+
+@pytest.mark.parametrize("upload_file_str", ["upload_file_ndjson", "upload_file"])
+def test_astra_stager(
+    request: TopRequest,
+    upload_file_str: str,
+    tmp_path: Path,
+):
+    upload_file: Path = request.getfixturevalue(upload_file_str)
+    stager = AstraDBUploadStager()
+    stager_validation(
+        configs=StagerValidationConfigs(test_id=CONNECTOR_TYPE, expected_count=22),
+        input_file=upload_file,
+        stager=stager,
+        tmp_dir=tmp_path,
+    )

test/integration/connectors/test_azure_ai_search.py

@@ -5,6 +5,7 @@ from pathlib import Path
 from uuid import uuid4
 
 import pytest
+from _pytest.fixtures import TopRequest
 from azure.core.credentials import AzureKeyCredential
 from azure.search.documents import SearchClient
 from azure.search.documents.indexes import SearchIndexClient
@@ -25,6 +26,10 @@ from azure.search.documents.indexes.models import (
 from test.integration.connectors.utils.constants import (
     DESTINATION_TAG,
 )
+from test.integration.connectors.utils.validation.destination import (
+    StagerValidationConfigs,
+    stager_validation,
+)
 from test.integration.utils import requires_env
 from unstructured_ingest.v2.interfaces import FileData, SourceIdentifiers
 from unstructured_ingest.v2.processes.connectors.azure_ai_search import (
@@ -225,9 +230,26 @@ async def test_azure_ai_search_destination(
     with staged_filepath.open() as f:
         staged_elements = json.load(f)
     expected_count = len(staged_elements)
-    search_client: SearchClient = uploader.connection_config.get_search_client()
-    validate_count(search_client=search_client, expected_count=expected_count)
+    with uploader.connection_config.get_search_client() as search_client:
+        validate_count(search_client=search_client, expected_count=expected_count)
 
     # Rerun and make sure the same documents get updated
     uploader.run(path=staged_filepath, file_data=file_data)
-    validate_count(search_client=search_client, expected_count=expected_count)
+    with uploader.connection_config.get_search_client() as search_client:
+        validate_count(search_client=search_client, expected_count=expected_count)
+
+
+@pytest.mark.parametrize("upload_file_str", ["upload_file_ndjson", "upload_file"])
+def test_azure_ai_search_stager(
+    request: TopRequest,
+    upload_file_str: str,
+    tmp_path: Path,
+):
+    upload_file: Path = request.getfixturevalue(upload_file_str)
+    stager = AzureAISearchUploadStager()
+    stager_validation(
+        configs=StagerValidationConfigs(test_id=CONNECTOR_TYPE, expected_count=22),
+        input_file=upload_file,
+        stager=stager,
+        tmp_dir=tmp_path,
+    )

test/integration/connectors/test_chroma.py

@@ -0,0 +1,120 @@
+import json
+from pathlib import Path
+
+import chromadb
+import pytest
+from _pytest.fixtures import TopRequest
+
+from test.integration.connectors.utils.constants import (
+    DESTINATION_TAG,
+)
+from test.integration.connectors.utils.docker import HealthCheck, container_context
+from test.integration.connectors.utils.validation.destination import (
+    StagerValidationConfigs,
+    stager_validation,
+)
+from unstructured_ingest.v2.interfaces import FileData, SourceIdentifiers
+from unstructured_ingest.v2.processes.connectors.chroma import (
+    CONNECTOR_TYPE,
+    ChromaConnectionConfig,
+    ChromaUploader,
+    ChromaUploaderConfig,
+    ChromaUploadStager,
+    ChromaUploadStagerConfig,
+)
+
+
+@pytest.fixture
+def chroma_instance():
+    with container_context(
+        image="chromadb/chroma:latest",
+        ports={8000: 8000},
+        name="chroma_int_test",
+        healthcheck=HealthCheck(
+            interval=5,
+            timeout=10,
+            retries=3,
+            test="timeout 10s bash -c ':> /dev/tcp/127.0.0.1/8000' || exit 1",
+        ),
+    ) as ctx:
+        yield ctx
+
+
+def validate_collection(collection_name: str, num_embeddings: int):
+    print(f"Checking contents of Chroma collection: {collection_name}")
+
+    chroma_client = chromadb.HttpClient(
+        host="localhost",
+        port="8000",
+        tenant="default_tenant",
+        database="default_database",
+    )
+
+    collection = chroma_client.get_or_create_collection(name=collection_name)
+
+    number_of_embeddings = collection.count()
+    expected_embeddings = num_embeddings
+    print(
+        f"# of embeddings in collection vs expected: {number_of_embeddings}/{expected_embeddings}"
+    )
+
+    assert number_of_embeddings == expected_embeddings, (
+        f"Number of rows in generated table ({number_of_embeddings}) "
+        f"doesn't match expected value: {expected_embeddings}"
+    )
+
+
+@pytest.mark.tags(CONNECTOR_TYPE, DESTINATION_TAG)
+def test_chroma_destination(
+    upload_file: Path,
+    chroma_instance,
+    tmp_path: Path,
+):
+    collection_name = "test_collection"
+    file_data = FileData(
+        source_identifiers=SourceIdentifiers(fullpath=upload_file.name, filename=upload_file.name),
+        connector_type=CONNECTOR_TYPE,
+        identifier="mock file data",
+    )
+    stager = ChromaUploadStager(upload_stager_config=ChromaUploadStagerConfig())
+
+    uploader = ChromaUploader(
+        connection_config=ChromaConnectionConfig(
+            host="localhost",
+            port=8000,
+            tenant="default_tenant",
+            database="default_database",
+        ),
+        upload_config=ChromaUploaderConfig(collection_name=collection_name),
+    )
+    staged_filepath = stager.run(
+        elements_filepath=upload_file,
+        file_data=file_data,
+        output_dir=tmp_path,
+        output_filename=upload_file.name,
+    )
+    uploader.precheck()
+    uploader.run(path=staged_filepath, file_data=file_data)
+
+    # Run validation
+    with staged_filepath.open() as f:
+        staged_elements = json.load(f)
+    expected_count = len(staged_elements)
+    validate_collection(collection_name=collection_name, num_embeddings=expected_count)
+
+
+@pytest.mark.parametrize("upload_file_str", ["upload_file_ndjson", "upload_file"])
+@pytest.mark.tags(CONNECTOR_TYPE, DESTINATION_TAG, "stager")
+def test_chroma_stager(
+    request: TopRequest,
+    upload_file_str: str,
+    tmp_path: Path,
+):
+    upload_file: Path = request.getfixturevalue(upload_file_str)
+    stager = ChromaUploadStager()
+    stager_validation(
+        configs=StagerValidationConfigs(test_id=CONNECTOR_TYPE, expected_count=22),
+        input_file=upload_file,
+        stager=stager,
+        tmp_dir=tmp_path,
+    )

test/integration/connectors/test_confluence.py

@@ -5,8 +5,8 @@ import pytest
 from test.integration.connectors.utils.constants import (
     SOURCE_TAG,
 )
-from test.integration.connectors.utils.validation import (
-    ValidationConfigs,
+from test.integration.connectors.utils.validation.source import (
+    SourceValidationConfigs,
     source_connector_validation,
 )
 from test.integration.utils import requires_env
@@ -60,7 +60,7 @@ async def test_confluence_source(temp_dir):
     await source_connector_validation(
         indexer=indexer,
         downloader=downloader,
-        configs=ValidationConfigs(
+        configs=SourceValidationConfigs(
             test_id="confluence",
             expected_num_files=11,
             validate_downloaded_files=True,
@@ -107,7 +107,7 @@ async def test_confluence_source_large(temp_dir):
     await source_connector_validation(
         indexer=indexer,
         downloader=downloader,
-        configs=ValidationConfigs(
+        configs=SourceValidationConfigs(
             test_id="confluence_large", expected_num_files=250, validate_file_data=False
         ),
     )

test/integration/connectors/test_delta_table.py

@@ -114,6 +114,7 @@ async def test_delta_table_destination_s3(upload_file: Path, temp_dir: Path):
     )
 
     try:
+        uploader.precheck()
         if uploader.is_async():
             await uploader.run_async(path=new_upload_file, file_data=file_data)
         else:

test/integration/connectors/test_kafka.py

@@ -14,8 +14,8 @@ from test.integration.connectors.utils.constants import (
     env_setup_path,
 )
 from test.integration.connectors.utils.docker_compose import docker_compose_context
-from test.integration.connectors.utils.validation import (
-    ValidationConfigs,
+from test.integration.connectors.utils.validation.source import (
+    SourceValidationConfigs,
     source_connector_validation,
 )
 from test.integration.utils import requires_env
@@ -121,7 +121,7 @@ async def test_kafka_source_local(kafka_seed_topic: str):
         await source_connector_validation(
             indexer=indexer,
             downloader=downloader,
-            configs=ValidationConfigs(
+            configs=SourceValidationConfigs(
                 test_id="kafka", expected_num_files=5, validate_downloaded_files=True
             ),
         )
@@ -203,7 +203,7 @@ async def test_kafka_source_cloud(kafka_seed_topic_cloud: int):
         await source_connector_validation(
             indexer=indexer,
             downloader=downloader,
-            configs=ValidationConfigs(
+            configs=SourceValidationConfigs(
                 test_id="kafka",
                 exclude_fields_extend=["connector_type"],
                 expected_num_files=expected_messages,

test/integration/connectors/test_milvus.py

@@ -4,6 +4,7 @@ from pathlib import Path
 
 import docker
 import pytest
+from _pytest.fixtures import TopRequest
 from pymilvus import (
     CollectionSchema,
     DataType,
@@ -15,6 +16,10 @@ from pymilvus.milvus_client import IndexParams
 from test.integration.connectors.utils.constants import DESTINATION_TAG, env_setup_path
 from test.integration.connectors.utils.docker import healthcheck_wait
 from test.integration.connectors.utils.docker_compose import docker_compose_context
+from test.integration.connectors.utils.validation.destination import (
+    StagerValidationConfigs,
+    stager_validation,
+)
 from unstructured_ingest.error import DestinationConnectionError
 from unstructured_ingest.v2.interfaces import FileData, SourceIdentifiers
 from unstructured_ingest.v2.processes.connectors.milvus import (
@@ -167,3 +172,19 @@ def test_precheck_fails_on_nonexistent_collection(collection: str):
         match=f"Collection '{NONEXISTENT_COLLECTION_NAME}' does not exist",
     ):
         uploader.precheck()
+
+
+@pytest.mark.parametrize("upload_file_str", ["upload_file_ndjson", "upload_file"])
+def test_milvus_stager(
+    request: TopRequest,
+    upload_file_str: str,
+    tmp_path: Path,
+):
+    upload_file: Path = request.getfixturevalue(upload_file_str)
+    stager = MilvusUploadStager()
+    stager_validation(
+        configs=StagerValidationConfigs(test_id=CONNECTOR_TYPE, expected_count=22),
+        input_file=upload_file,
+        stager=stager,
+        tmp_dir=tmp_path,
+    )

test/integration/connectors/test_mongodb.py

@@ -14,8 +14,8 @@ from pymongo.mongo_client import MongoClient
 from pymongo.operations import SearchIndexModel
 
 from test.integration.connectors.utils.constants import DESTINATION_TAG, SOURCE_TAG
-from test.integration.connectors.utils.validation import (
-    ValidationConfigs,
+from test.integration.connectors.utils.validation.source import (
+    SourceValidationConfigs,
     source_connector_validation,
 )
 from test.integration.utils import requires_env
@@ -196,7 +196,7 @@ async def test_mongodb_source(temp_dir: Path):
     await source_connector_validation(
         indexer=indexer,
         downloader=downloader,
-        configs=ValidationConfigs(
+        configs=SourceValidationConfigs(
             test_id=CONNECTOR_TYPE, expected_num_files=4, validate_downloaded_files=True
         ),
     )