unstructured-ingest 0.3.0__py3-none-any.whl → 0.3.2__py3-none-any.whl

test/integration/connectors/sql/test_snowflake.py

@@ -86,7 +86,7 @@ async def test_snowflake_source():
         image="localstack/snowflake",
         environment={"LOCALSTACK_AUTH_TOKEN": token, "EXTRA_CORS_ALLOWED_ORIGINS": "*"},
         ports={4566: 4566, 443: 443},
-        healthcheck_timeout=30,
+        healthcheck_retries=30,
     ):
         seed_data()
         with tempfile.TemporaryDirectory() as tmpdir:
@@ -156,7 +156,7 @@ async def test_snowflake_destination(upload_file: Path):
         image="localstack/snowflake",
         environment={"LOCALSTACK_AUTH_TOKEN": token, "EXTRA_CORS_ALLOWED_ORIGINS": "*"},
         ports={4566: 4566, 443: 443},
-        healthcheck_timeout=30,
+        healthcheck_retries=30,
     ):
         init_db_destination()
         with tempfile.TemporaryDirectory() as tmpdir:
@@ -192,10 +192,8 @@ async def test_snowflake_destination(upload_file: Path):
                     host=connect_params["host"],
                 )
             )
-            if uploader.is_async():
-                await uploader.run_async(path=staged_path, file_data=mock_file_data)
-            else:
-                uploader.run(path=staged_path, file_data=mock_file_data)
+
+            uploader.run(path=staged_path, file_data=mock_file_data)
 
             staged_df = pd.read_json(staged_path, orient="records", lines=True)
             expected_num_elements = len(staged_df)
@@ -203,3 +201,9 @@ async def test_snowflake_destination(upload_file: Path):
                 connect_params=connect_params,
                 expected_num_elements=expected_num_elements,
             )
+
+            uploader.run(path=staged_path, file_data=mock_file_data)
+            validate_destination(
+                connect_params=connect_params,
+                expected_num_elements=expected_num_elements,
+            )

test/integration/connectors/sql/test_sqlite.py

@@ -138,10 +138,10 @@ async def test_sqlite_destination(upload_file: Path):
             uploader = SQLiteUploader(
                 connection_config=SQLiteConnectionConfig(database_path=db_path)
             )
-            if uploader.is_async():
-                await uploader.run_async(path=staged_path, file_data=mock_file_data)
-            else:
-                uploader.run(path=staged_path, file_data=mock_file_data)
+            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))

test/integration/connectors/test_astradb.py

@@ -8,20 +8,67 @@ import pytest
 from astrapy import Collection
 from astrapy import DataAPIClient as AstraDBClient
 
-from test.integration.connectors.utils.constants import (
-    DESTINATION_TAG,
-)
+from test.integration.connectors.utils.constants import DESTINATION_TAG, SOURCE_TAG
 from test.integration.utils import requires_env
 from unstructured_ingest.v2.interfaces import FileData, SourceIdentifiers
 from unstructured_ingest.v2.processes.connectors.astradb import (
     CONNECTOR_TYPE,
     AstraDBAccessConfig,
     AstraDBConnectionConfig,
+    AstraDBIndexer,
+    AstraDBIndexerConfig,
     AstraDBUploader,
     AstraDBUploaderConfig,
     AstraDBUploadStager,
+    DestinationConnectionError,
+    SourceConnectionError,
 )
 
+EXISTENT_COLLECTION_NAME = "ingest_test_src"
+NONEXISTENT_COLLECTION_NAME = "nonexistant"
+
+
+@pytest.fixture
+def connection_config() -> AstraDBConnectionConfig:
+    return AstraDBConnectionConfig(
+        access_config=AstraDBAccessConfig(
+            token=os.environ["ASTRA_DB_APPLICATION_TOKEN"],
+            api_endpoint=os.environ["ASTRA_DB_API_ENDPOINT"],
+        )
+    )
+
+
+@pytest.mark.tags(CONNECTOR_TYPE, SOURCE_TAG, DESTINATION_TAG)
+@requires_env("ASTRA_DB_APPLICATION_TOKEN", "ASTRA_DB_API_ENDPOINT")
+def test_precheck_succeeds(connection_config: AstraDBConnectionConfig):
+    indexer = AstraDBIndexer(
+        connection_config=connection_config,
+        index_config=AstraDBIndexerConfig(collection_name=EXISTENT_COLLECTION_NAME),
+    )
+    uploader = AstraDBUploader(
+        connection_config=connection_config,
+        upload_config=AstraDBUploaderConfig(collection_name=EXISTENT_COLLECTION_NAME),
+    )
+    indexer.precheck()
+    uploader.precheck()
+
+
+@pytest.mark.tags(CONNECTOR_TYPE, SOURCE_TAG, DESTINATION_TAG)
+@requires_env("ASTRA_DB_APPLICATION_TOKEN", "ASTRA_DB_API_ENDPOINT")
+def test_precheck_fails(connection_config: AstraDBConnectionConfig):
+    indexer = AstraDBIndexer(
+        connection_config=connection_config,
+        index_config=AstraDBIndexerConfig(collection_name=NONEXISTENT_COLLECTION_NAME),
+    )
+    uploader = AstraDBUploader(
+        connection_config=connection_config,
+        upload_config=AstraDBUploaderConfig(collection_name=NONEXISTENT_COLLECTION_NAME),
+    )
+    with pytest.raises(expected_exception=SourceConnectionError):
+        indexer.precheck()
+    with pytest.raises(expected_exception=DestinationConnectionError):
+        uploader.precheck()
+
 
 @dataclass(frozen=True)
 class EnvData:

test/integration/connectors/test_delta_table.py

@@ -136,3 +136,49 @@ async def test_delta_table_destination_s3(upload_file: Path, temp_dir: Path):
             secret=aws_credentials["AWS_SECRET_ACCESS_KEY"],
         )
         s3fs.rm(path=destination_path, recursive=True)
+
+
+@pytest.mark.asyncio
+@pytest.mark.tags(CONNECTOR_TYPE, DESTINATION_TAG)
+@requires_env("S3_INGEST_TEST_ACCESS_KEY", "S3_INGEST_TEST_SECRET_KEY")
+async def test_delta_table_destination_s3_bad_creds(upload_file: Path, temp_dir: Path):
+    aws_credentials = {
+        "AWS_ACCESS_KEY_ID": "bad key",
+        "AWS_SECRET_ACCESS_KEY": "bad secret",
+        "AWS_REGION": "us-east-2",
+    }
+    s3_bucket = "s3://utic-platform-test-destination"
+    destination_path = f"{s3_bucket}/destination/test"
+    connection_config = DeltaTableConnectionConfig(
+        access_config=DeltaTableAccessConfig(
+            aws_access_key_id=aws_credentials["AWS_ACCESS_KEY_ID"],
+            aws_secret_access_key=aws_credentials["AWS_SECRET_ACCESS_KEY"],
+        ),
+        aws_region=aws_credentials["AWS_REGION"],
+        table_uri=destination_path,
+    )
+    stager_config = DeltaTableUploadStagerConfig()
+    stager = DeltaTableUploadStager(upload_stager_config=stager_config)
+    new_upload_file = stager.run(
+        elements_filepath=upload_file,
+        output_dir=temp_dir,
+        output_filename=upload_file.name,
+    )
+
+    upload_config = DeltaTableUploaderConfig()
+    uploader = DeltaTableUploader(connection_config=connection_config, upload_config=upload_config)
+    file_data = FileData(
+        source_identifiers=SourceIdentifiers(
+            fullpath=upload_file.name, filename=new_upload_file.name
+        ),
+        connector_type=CONNECTOR_TYPE,
+        identifier="mock file data",
+    )
+
+    with pytest.raises(Exception) as excinfo:
+        if uploader.is_async():
+            await uploader.run_async(path=new_upload_file, file_data=file_data)
+        else:
+            uploader.run(path=new_upload_file, file_data=file_data)
+
+    assert "403 Forbidden" in str(excinfo.value), f"Exception message did not match: {str(excinfo)}"

test/integration/connectors/test_kafka.py

@@ -1,5 +1,6 @@
 import json
 import tempfile
+import time
 from pathlib import Path
 
 import pytest
@@ -33,12 +34,36 @@ SEED_MESSAGES = 10
 TOPIC = "fake-topic"
 
 
+def get_admin_client() -> AdminClient:
+    conf = {
+        "bootstrap.servers": "localhost:29092",
+    }
+    return AdminClient(conf)
+
+
 @pytest.fixture
 def docker_compose_ctx():
     with docker_compose_context(docker_compose_path=env_setup_path / "kafka") as ctx:
         yield ctx
 
 
+def wait_for_topic(topic: str, retries: int = 10, interval: int = 1):
+    admin_client = get_admin_client()
+    current_topics = admin_client.list_topics().topics
+    attempts = 0
+    while topic not in current_topics and attempts < retries:
+        attempts += 1
+        print(
+            "Attempt {}: Waiting for topic {} to exist in {}".format(
+                attempts, topic, ", ".join(current_topics)
+            )
+        )
+        time.sleep(interval)
+        current_topics = admin_client.list_topics().topics
+    if topic not in current_topics:
+        raise TimeoutError(f"Timeout out waiting for topic {topic} to exist")
+
+
 @pytest.fixture
 def kafka_seed_topic(docker_compose_ctx) -> str:
     conf = {
@@ -50,15 +75,13 @@ def kafka_seed_topic(docker_compose_ctx) -> str:
         producer.produce(topic=TOPIC, value=message)
     producer.flush(timeout=10)
     print(f"kafka topic {TOPIC} seeded with {SEED_MESSAGES} messages")
+    wait_for_topic(topic=TOPIC)
     return TOPIC
 
 
 @pytest.fixture
 def kafka_upload_topic(docker_compose_ctx) -> str:
-    conf = {
-        "bootstrap.servers": "localhost:29092",
-    }
-    admin_client = AdminClient(conf)
+    admin_client = get_admin_client()
     admin_client.create_topics([NewTopic(TOPIC, 1, 1)])
     return TOPIC
 
@@ -88,7 +111,7 @@ async def test_kafka_source_local(kafka_seed_topic: str):
 
 
 @pytest.mark.tags(CONNECTOR_TYPE, SOURCE_TAG)
-def test_kafak_source_local_precheck_fail():
+def test_kafka_source_local_precheck_fail_no_cluster():
     connection_config = LocalKafkaConnectionConfig(bootstrap_server="localhost", port=29092)
     indexer = LocalKafkaIndexer(
         connection_config=connection_config,
@@ -98,6 +121,17 @@ def test_kafak_source_local_precheck_fail():
         indexer.precheck()
 
 
+@pytest.mark.tags(CONNECTOR_TYPE, SOURCE_TAG)
+def test_kafka_source_local_precheck_fail_no_topic(kafka_seed_topic: str):
+    connection_config = LocalKafkaConnectionConfig(bootstrap_server="localhost", port=29092)
+    indexer = LocalKafkaIndexer(
+        connection_config=connection_config,
+        index_config=LocalKafkaIndexerConfig(topic="topic", num_messages_to_consume=5),
+    )
+    with pytest.raises(SourceConnectionError):
+        indexer.precheck()
+
+
 def get_all_messages(topic: str, max_empty_messages: int = 5) -> list[dict]:
     conf = {
         "bootstrap.servers": "localhost:29092",
@@ -158,7 +192,7 @@ async def test_kafka_destination_local(upload_file: Path, kafka_upload_topic: st
 
 
 @pytest.mark.tags(CONNECTOR_TYPE, DESTINATION_TAG)
-def test_kafak_destination_local_precheck_fail():
+def test_kafka_destination_local_precheck_fail_no_cluster():
     uploader = LocalKafkaUploader(
         connection_config=LocalKafkaConnectionConfig(bootstrap_server="localhost", port=29092),
         upload_config=LocalKafkaUploaderConfig(topic=TOPIC, batch_size=10),

test/integration/connectors/test_lancedb.py

@@ -0,0 +1,210 @@
+import os
+from pathlib import Path
+from typing import Literal, Union
+from uuid import uuid4
+
+import lancedb
+import pandas as pd
+import pyarrow as pa
+import pytest
+import pytest_asyncio
+from lancedb import AsyncConnection
+from upath import UPath
+
+from test.integration.connectors.utils.constants import DESTINATION_TAG
+from unstructured_ingest.v2.interfaces.file_data import FileData, SourceIdentifiers
+from unstructured_ingest.v2.processes.connectors.lancedb.aws import (
+    LanceDBS3AccessConfig,
+    LanceDBS3ConnectionConfig,
+    LanceDBS3Uploader,
+)
+from unstructured_ingest.v2.processes.connectors.lancedb.azure import (
+    LanceDBAzureAccessConfig,
+    LanceDBAzureConnectionConfig,
+    LanceDBAzureUploader,
+)
+from unstructured_ingest.v2.processes.connectors.lancedb.gcp import (
+    LanceDBGCSAccessConfig,
+    LanceDBGCSConnectionConfig,
+    LanceDBGSPUploader,
+)
+from unstructured_ingest.v2.processes.connectors.lancedb.lancedb import (
+    CONNECTOR_TYPE,
+    LanceDBUploaderConfig,
+    LanceDBUploadStager,
+)
+from unstructured_ingest.v2.processes.connectors.lancedb.local import (
+    LanceDBLocalAccessConfig,
+    LanceDBLocalConnectionConfig,
+    LanceDBLocalUploader,
+)
+
+DATABASE_NAME = "database"
+TABLE_NAME = "elements"
+DIMENSION = 384
+NUMBER_EXPECTED_ROWS = 22
+NUMBER_EXPECTED_COLUMNS = 10
+S3_BUCKET = "s3://utic-ingest-test-fixtures/"
+GS_BUCKET = "gs://utic-test-ingest-fixtures-output/"
+AZURE_BUCKET = "az://utic-ingest-test-fixtures-output/"
+REQUIRED_ENV_VARS = {
+    "s3": ("S3_INGEST_TEST_ACCESS_KEY", "S3_INGEST_TEST_SECRET_KEY"),
+    "gcs": ("GCP_INGEST_SERVICE_KEY",),
+    "az": ("AZURE_DEST_CONNECTION_STR",),
+    "local": (),
+}
+
+
+SCHEMA = pa.schema(
+    [
+        pa.field("vector", pa.list_(pa.float16(), DIMENSION)),
+        pa.field("text", pa.string(), nullable=True),
+        pa.field("type", pa.string(), nullable=True),
+        pa.field("element_id", pa.string(), nullable=True),
+        pa.field("metadata-text_as_html", pa.string(), nullable=True),
+        pa.field("metadata-filetype", pa.string(), nullable=True),
+        pa.field("metadata-filename", pa.string(), nullable=True),
+        pa.field("metadata-languages", pa.list_(pa.string()), nullable=True),
+        pa.field("metadata-is_continuation", pa.bool_(), nullable=True),
+        pa.field("metadata-page_number", pa.int32(), nullable=True),
+    ]
+)
+
+
+@pytest_asyncio.fixture
+async def connection_with_uri(request, tmp_path: Path):
+    target = request.param
+    uri = _get_uri(target, local_base_path=tmp_path)
+
+    unset_variables = [env for env in REQUIRED_ENV_VARS[target] if env not in os.environ]
+    if unset_variables:
+        pytest.skip(
+            reason="Following required environment variables were not set: "
+            + f"{', '.join(unset_variables)}"
+        )
+
+    storage_options = {
+        "aws_access_key_id": os.getenv("S3_INGEST_TEST_ACCESS_KEY"),
+        "aws_secret_access_key": os.getenv("S3_INGEST_TEST_SECRET_KEY"),
+        "google_service_account_key": os.getenv("GCP_INGEST_SERVICE_KEY"),
+    }
+    azure_connection_string = os.getenv("AZURE_DEST_CONNECTION_STR")
+    if azure_connection_string:
+        storage_options.update(_parse_azure_connection_string(azure_connection_string))
+
+    storage_options = {key: value for key, value in storage_options.items() if value is not None}
+    connection = await lancedb.connect_async(
+        uri=uri,
+        storage_options=storage_options,
+    )
+    await connection.create_table(name=TABLE_NAME, schema=SCHEMA)
+
+    yield connection, uri
+
+    await connection.drop_database()
+
+
+@pytest.mark.asyncio
+@pytest.mark.tags(CONNECTOR_TYPE, DESTINATION_TAG)
+@pytest.mark.parametrize("connection_with_uri", ["local", "s3", "gcs", "az"], indirect=True)
+async def test_lancedb_destination(
+    upload_file: Path,
+    connection_with_uri: tuple[AsyncConnection, str],
+    tmp_path: Path,
+) -> None:
+    connection, uri = connection_with_uri
+    file_data = FileData(
+        source_identifiers=SourceIdentifiers(fullpath=upload_file.name, filename=upload_file.name),
+        connector_type=CONNECTOR_TYPE,
+        identifier="mock file data",
+    )
+    stager = LanceDBUploadStager()
+    uploader = _get_uploader(uri)
+    staged_file_path = stager.run(
+        elements_filepath=upload_file,
+        file_data=file_data,
+        output_dir=tmp_path,
+        output_filename=upload_file.name,
+    )
+
+    await uploader.run_async(path=staged_file_path, file_data=file_data)
+
+    table = await connection.open_table(TABLE_NAME)
+    table_df: pd.DataFrame = await table.to_pandas()
+
+    assert len(table_df) == NUMBER_EXPECTED_ROWS
+    assert len(table_df.columns) == NUMBER_EXPECTED_COLUMNS
+
+    assert table_df["element_id"][0] == "2470d8dc42215b3d68413b55bf00fed2"
+    assert table_df["type"][0] == "CompositeElement"
+    assert table_df["metadata-filename"][0] == "DA-1p-with-duplicate-pages.pdf.json"
+    assert table_df["metadata-text_as_html"][0] is None
+
+
+def _get_uri(target: Literal["local", "s3", "gcs", "az"], local_base_path: Path) -> str:
+    if target == "local":
+        return str(local_base_path / DATABASE_NAME)
+    if target == "s3":
+        base_uri = UPath(S3_BUCKET)
+    elif target == "gcs":
+        base_uri = UPath(GS_BUCKET)
+    elif target == "az":
+        base_uri = UPath(AZURE_BUCKET)
+
+    return str(base_uri / "destination" / "lancedb" / str(uuid4()) / DATABASE_NAME)
+
+
+def _get_uploader(
+    uri: str,
+) -> Union[LanceDBAzureUploader, LanceDBAzureUploader, LanceDBS3Uploader, LanceDBGSPUploader]:
+    target = uri.split("://", maxsplit=1)[0] if uri.startswith(("s3", "az", "gs")) else "local"
+    if target == "az":
+        azure_connection_string = os.getenv("AZURE_DEST_CONNECTION_STR")
+        access_config_kwargs = _parse_azure_connection_string(azure_connection_string)
+        return LanceDBAzureUploader(
+            upload_config=LanceDBUploaderConfig(table_name=TABLE_NAME),
+            connection_config=LanceDBAzureConnectionConfig(
+                access_config=LanceDBAzureAccessConfig(**access_config_kwargs),
+                uri=uri,
+            ),
+        )
+
+    elif target == "s3":
+        return LanceDBS3Uploader(
+            upload_config=LanceDBUploaderConfig(table_name=TABLE_NAME),
+            connection_config=LanceDBS3ConnectionConfig(
+                access_config=LanceDBS3AccessConfig(
+                    aws_access_key_id=os.getenv("S3_INGEST_TEST_ACCESS_KEY"),
+                    aws_secret_access_key=os.getenv("S3_INGEST_TEST_SECRET_KEY"),
+                ),
+                uri=uri,
+            ),
+        )
+    elif target == "gs":
+        return LanceDBGSPUploader(
+            upload_config=LanceDBUploaderConfig(table_name=TABLE_NAME),
+            connection_config=LanceDBGCSConnectionConfig(
+                access_config=LanceDBGCSAccessConfig(
+                    google_service_account_key=os.getenv("GCP_INGEST_SERVICE_KEY")
+                ),
+                uri=uri,
+            ),
+        )
+    else:
+        return LanceDBLocalUploader(
+            upload_config=LanceDBUploaderConfig(table_name=TABLE_NAME),
+            connection_config=LanceDBLocalConnectionConfig(
+                access_config=LanceDBLocalAccessConfig(),
+                uri=uri,
+            ),
+        )
+
+
+def _parse_azure_connection_string(
+    connection_str: str,
+) -> dict[Literal["azure_storage_account_name", "azure_storage_account_key"], str]:
+    parameters = dict(keyvalue.split("=", maxsplit=1) for keyvalue in connection_str.split(";"))
+    return {
+        "azure_storage_account_name": parameters.get("AccountName"),
+        "azure_storage_account_key": parameters.get("AccountKey"),
+    }

test/integration/connectors/test_milvus.py

@@ -0,0 +1,141 @@
+import json
+import time
+from pathlib import Path
+
+import docker
+import pytest
+from pymilvus import (
+    CollectionSchema,
+    DataType,
+    FieldSchema,
+    MilvusClient,
+)
+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 unstructured_ingest.v2.interfaces import FileData, SourceIdentifiers
+from unstructured_ingest.v2.processes.connectors.milvus import (
+    CONNECTOR_TYPE,
+    MilvusConnectionConfig,
+    MilvusUploader,
+    MilvusUploaderConfig,
+    MilvusUploadStager,
+)
+
+DB_URI = "http://localhost:19530"
+DB_NAME = "test_database"
+COLLECTION_NAME = "test_collection"
+
+
+def get_schema() -> CollectionSchema:
+    id_field = FieldSchema(
+        name="id", dtype=DataType.INT64, description="primary field", is_primary=True, auto_id=True
+    )
+    embeddings_field = FieldSchema(name="embeddings", dtype=DataType.FLOAT_VECTOR, dim=384)
+    record_id_field = FieldSchema(name="record_id", dtype=DataType.VARCHAR, max_length=64)
+
+    schema = CollectionSchema(
+        enable_dynamic_field=True,
+        fields=[
+            id_field,
+            record_id_field,
+            embeddings_field,
+        ],
+    )
+
+    return schema
+
+
+def get_index_params() -> IndexParams:
+    index_params = IndexParams()
+    index_params.add_index(field_name="embeddings", index_type="AUTOINDEX", metric_type="COSINE")
+    index_params.add_index(field_name="record_id", index_type="Trie")
+    return index_params
+
+
+@pytest.fixture
+def collection():
+    docker_client = docker.from_env()
+    with docker_compose_context(docker_compose_path=env_setup_path / "milvus"):
+        milvus_container = docker_client.containers.get("milvus-standalone")
+        healthcheck_wait(container=milvus_container)
+        milvus_client = MilvusClient(uri=DB_URI)
+        try:
+            # Create the database
+            database_resp = milvus_client._get_connection().create_database(db_name=DB_NAME)
+            milvus_client.using_database(db_name=DB_NAME)
+
+            print(f"Created database {DB_NAME}: {database_resp}")
+
+            # Create the collection
+            schema = get_schema()
+            index_params = get_index_params()
+            collection_resp = milvus_client.create_collection(
+                collection_name=COLLECTION_NAME, schema=schema, index_params=index_params
+            )
+            print(f"Created collection {COLLECTION_NAME}: {collection_resp}")
+            yield COLLECTION_NAME
+        finally:
+            milvus_client.close()
+
+
+def get_count(client: MilvusClient) -> int:
+    count_field = "count(*)"
+    resp = client.query(collection_name="test_collection", output_fields=[count_field])
+    return resp[0][count_field]
+
+
+def validate_count(
+    client: MilvusClient, expected_count: int, retries: int = 10, interval: int = 1
+) -> None:
+    current_count = get_count(client=client)
+    retry_count = 0
+    while current_count != expected_count and retry_count < retries:
+        time.sleep(interval)
+        current_count = get_count(client=client)
+        retry_count += 1
+    assert current_count == expected_count, (
+        f"Expected count ({expected_count}) doesn't match how "
+        f"much came back from collection: {current_count}"
+    )
+
+
+@pytest.mark.asyncio
+@pytest.mark.tags(CONNECTOR_TYPE, DESTINATION_TAG)
+async def test_milvus_destination(
+    upload_file: Path,
+    collection: str,
+    tmp_path: Path,
+):
+    file_data = FileData(
+        source_identifiers=SourceIdentifiers(fullpath=upload_file.name, filename=upload_file.name),
+        connector_type=CONNECTOR_TYPE,
+        identifier="mock file data",
+    )
+    stager = MilvusUploadStager()
+    uploader = MilvusUploader(
+        connection_config=MilvusConnectionConfig(uri=DB_URI),
+        upload_config=MilvusUploaderConfig(collection_name=collection, db_name=DB_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)
+    with uploader.get_client() as client:
+        validate_count(client=client, expected_count=expected_count)
+
+    # Rerun and make sure the same documents get updated
+    uploader.run(path=staged_filepath, file_data=file_data)
+    with uploader.get_client() as client:
+        validate_count(client=client, expected_count=expected_count)