apache_beam.ml.rag.test_utils module

class apache_beam.ml.rag.test_utils.VectorDBContainerInfo(container: DbContainer, host: str, port: int, user: str = '', password: str = '', token: str = '', id: str = 'default')

Bases: object

Container information for vector database test instances.

Holds connection details and container reference for testing with vector databases like Milvus in containerized environments.

container: DbContainer

host: str

port: int

user: str = ''

password: str = ''

token: str = ''

id: str = 'default'

property uri: str

class apache_beam.ml.rag.test_utils.TestHelpers

Bases: object

static find_free_port()

Find a free port on the local machine.

class apache_beam.ml.rag.test_utils.CustomMilvusContainer(image: str, service_container_port, healthcheck_container_port, **kwargs)

Bases: MilvusContainer

Custom Milvus container with configurable ports and environment setup.

Extends MilvusContainer to provide custom port binding and environment configuration for testing with standalone Milvus instances.

class apache_beam.ml.rag.test_utils.MilvusTestHelpers

Bases: object

Helper utilities for testing Milvus vector database operations.

Provides static methods for managing test containers, configuration files, and chunk comparison utilities for Milvus-based integration tests.

static start_db_container(image='milvusdb/milvus:v2.5.10', max_vec_fields=5, vector_client_max_retries=3, tc_max_retries=None) → VectorDBContainerInfo | None

static stop_db_container(db_info: VectorDBContainerInfo)

static initialize_db_with_data(connc_params: MilvusConnectionParameters, config: dict)

static create_user_yaml(service_port: int, max_vector_field_num=5)

Creates a temporary user.yaml file for Milvus configuration.

This user yaml file overrides Milvus default configurations. It sets the Milvus service port to the specified container service port. The default for maxVectorFieldNum is 4, but we need 5 (one unique field for each metric).

Parameters:

• service_port – Port number for the Milvus service.

• max_vector_field_num – Max number of vec fields allowed per collection.

Yields:

str – Path to the created temporary yaml file.

static assert_chunks_equivalent(actual_chunks: List[Chunk], expected_chunks: List[Chunk])

assert_chunks_equivalent checks for presence rather than exact match