"""Helper functions for pytest fixtures and container management.""" import gzip import os import shutil import socket import ssl import subprocess import tarfile import time import uuid import zipfile from pathlib import Path from typing import Dict, Any, Tuple, Optional, Union, List import docker.models.containers import requests from docker.errors import NotFound from .models import QdrantContainer, QdrantContainerConfig, QdrantDockerCluster def remove_dir(path: Path) -> None: """Remove a directory and all its contents. Args: path: Path to the directory to remove """ if path.exists(): shutil.rmtree(path, ignore_errors=True) print(f"Removed directory: {path}") def wait_for_qdrant_ready( port: int = 6333, timeout: int = 30, *, host: str = "localhost", scheme: str = "http", verify: Union[bool, str] = True, cert: Optional[Tuple[str, str]] = None, grpc_port: Optional[int] = None, ) -> bool: """Poll /readyz; if grpc_port is set, also require a TCP/TLS handshake on it.""" url = f"{scheme}://{host}:{port}/readyz" tls_ctx: Optional[ssl.SSLContext] = None if grpc_port is not None and scheme == "https": tls_ctx = ssl.create_default_context(cafile=verify if isinstance(verify, str) else None) if cert is not None: tls_ctx.load_cert_chain(certfile=cert[0], keyfile=cert[1]) deadline = time.time() + timeout while time.time() < deadline: try: if (requests.get(url, timeout=5, verify=verify, cert=cert).status_code == 200 and _grpc_port_ready(host, grpc_port, tls_ctx)): return True except requests.exceptions.RequestException: pass time.sleep(0.2) return False def _grpc_port_ready(host: str, grpc_port: Optional[int], tls_ctx: Optional[ssl.SSLContext]) -> bool: if grpc_port is None: return True try: with socket.create_connection((host, grpc_port), timeout=5) as sock: if tls_ctx is None: return True with tls_ctx.wrap_socket(sock, server_hostname=host): return True except (OSError, ssl.SSLError): return False def get_docker_compose_command() -> List[str]: """Detect and return the available docker-compose command. Tries docker compose v2 first, then falls back to docker-compose v1. Returns: List[str]: Command prefix for docker-compose (e.g., ["docker", "compose"] or ["docker-compose"]) Raises: RuntimeError: If neither docker compose nor docker-compose is available """ compose_commands = [ ["docker", "compose"], # v2 ["docker-compose"] # v1 ] for cmd_prefix in compose_commands: test_cmd = cmd_prefix + ["version"] result = subprocess.run(test_cmd, capture_output=True, text=True) if result.returncode == 0: return cmd_prefix raise RuntimeError("Neither 'docker compose' nor 'docker-compose' command found") def get_default_qdrant_config(qdrant_image: str) -> Dict[str, Any]: """Get default configuration for Qdrant container. Args: qdrant_image: The Qdrant Docker image to use Returns: dict: Default container configuration with image, ports, detach, and remove settings """ return { "image": qdrant_image, "ports": {'6333/tcp': ('127.0.0.1', None), '6334/tcp': ('127.0.0.1', None)}, "detach": True, "remove": True, } def extract_container_ports(container: docker.models.containers.Container) -> Tuple[int, int]: """Extract HTTP and gRPC ports from container. For host network mode, returns standard Qdrant ports (6333, 6334). For bridge/custom networks, extracts mapped ports from container attributes. Args: container: Docker container object Returns: tuple: (http_port, grpc_port) - extracted port numbers Raises: NotFound: If container has already been removed (e.g., crashed with auto-remove) RuntimeError: If port mappings cannot be extracted """ container.reload() if container.attrs.get('HostConfig', {}).get('NetworkMode') == 'host': # For host network mode, use standard Qdrant ports return 6333, 6334 # For bridge/custom networks, extract mapped ports ports = container.attrs.get('NetworkSettings', {}).get('Ports', {}) http_mapping = ports.get('6333/tcp') grpc_mapping = ports.get('6334/tcp') if not http_mapping: raise RuntimeError(f"Container {container.name} has no HTTP port mapping. Container may have exited.") http_port = int(http_mapping[0]['HostPort']) grpc_port = int(grpc_mapping[0]['HostPort']) if grpc_mapping else 6334 return http_port, grpc_port def create_container_info(container: docker.models.containers.Container, http_port: int, grpc_port: int) -> QdrantContainer: """Create standardized container info object. Args: container: Docker container object http_port: HTTP API port number grpc_port: gRPC API port number Returns: QdrantContainer: Container info object with container, host, name, http_port, and grpc_port attributes """ container.reload() if container.attrs.get('HostConfig', {}).get('NetworkMode') == 'host': host = "localhost" else: host = "127.0.0.1" return QdrantContainer( container=container, host=host, name=container.name, http_port=http_port, grpc_port=grpc_port ) def cleanup_container(container: docker.models.containers.Container) -> None: """Clean up a Docker container. Stops the container and removes it if AutoRemove is not enabled. Handles NotFound exceptions gracefully. Args: container: Docker container object to clean up """ try: container.reload() container.stop() if not container.attrs.get('HostConfig', {}).get('AutoRemove', True): container.remove(force=True) print(f"Removed container: {container.name}") except NotFound: print("Container already removed. OK.") except Exception as e: print(f"Error stopping container {container.name if hasattr(container, 'name') else 'unknown'}: {e}") def create_qdrant_container(docker_client: docker.DockerClient, qdrant_image: str, config: Optional[Union[Dict[str, Any], QdrantContainerConfig]] = None) -> QdrantContainer: """Core function to create a Qdrant container with given configuration. Args: docker_client: Docker client instance qdrant_image: Qdrant Docker image to use config: Optional configuration (dict or QdrantContainerConfig). Special parameters: - exit_on_error (bool): If True (default), raises RuntimeError when Qdrant fails to start. If False, returns container info even if Qdrant doesn't start. All other parameters are passed to docker_client.containers.run() Returns: QdrantContainer: Container info object (see create_container_info) Raises: RuntimeError: If Qdrant fails to start and exit_on_error=True """ if config is None: config = {} # Handle both dict and QdrantContainerConfig inputs if isinstance(config, QdrantContainerConfig): exit_on_error = config.exit_on_error merged_config = config.to_docker_config(qdrant_image) else: config = dict(config) # Extract custom parameters exit_on_error = config.pop("exit_on_error", True) default_config = get_default_qdrant_config(qdrant_image) merged_config = {**default_config, **config} # If using host network mode, remove port bindings if merged_config.get('network_mode') == 'host': merged_config.pop('ports', None) container = docker_client.containers.run(**merged_config) try: http_port, grpc_port = extract_container_ports(container) if not wait_for_qdrant_ready(port=http_port, timeout=30): if exit_on_error: raise RuntimeError("Qdrant failed to start within 30 seconds") return create_container_info(container, http_port, grpc_port) except Exception: cleanup_container(container) raise def extract_archive(archive_file: Path, extract_to: Path, cleanup_archive: bool = False) -> Path: """General utility function to extract various archive formats. Args: archive_file: Path to the archive file to extract extract_to: Directory where to extract the archive contents cleanup_archive: Whether to delete the archive file after extraction Returns: Path: The extraction directory path Raises: FileNotFoundError: If archive file doesn't exist ValueError: If archive format is unsupported RuntimeError: If extraction fails Supported formats: .tar.xz, .tar.gz, .tar.bz2, .tgz, .tbz2, .tar, .zip, .gz """ if not archive_file.exists(): raise FileNotFoundError(f"Archive not found: {archive_file}") extract_to.mkdir(parents=True, exist_ok=True) file_name = archive_file.name.lower() try: if file_name.endswith('.gz') and not file_name.endswith(('.tar.gz', '.tgz')): # Handle standalone gzip files (like snapshots) output_file = extract_to / archive_file.stem with gzip.open(archive_file, 'rb') as gz_file, \ open(output_file, 'wb') as out_file: shutil.copyfileobj(gz_file, out_file) print(f"Extracted {archive_file} to {output_file}") elif file_name.endswith(('.tar.xz', '.tar.gz', '.tar.bz2', '.tgz', '.tbz2')): # Handle compressed tar files with tarfile.open(archive_file, 'r:*') as tar: tar.extractall(path=extract_to, filter='data') print(f"Extracted {archive_file} to {extract_to}") elif file_name.endswith('.tar'): # Handle uncompressed tar files with tarfile.open(archive_file, 'r:') as tar: tar.extractall(path=extract_to, filter='data') print(f"Extracted {archive_file} to {extract_to}") elif file_name.endswith('.zip'): # Handle zip files with zipfile.ZipFile(archive_file, 'r') as zip_file: zip_file.extractall(path=extract_to) print(f"Extracted {archive_file} to {extract_to}") else: raise ValueError(f"Unsupported archive format: {archive_file}") except (OSError, ValueError, tarfile.TarError, zipfile.BadZipFile) as e: print(f"Failed to extract archive {archive_file}: {e}") # Try fallback to subprocess for tar files if file_name.endswith(('.tar.xz', '.tar.gz', '.tar.bz2', '.tgz', '.tbz2', '.tar')): try: print("Trying fallback extraction with tar command...") subprocess.run(["tar", "-xf", str(archive_file)], cwd=str(extract_to), check=True) print(f"Successfully extracted {archive_file} using tar command") except subprocess.CalledProcessError as tar_error: raise RuntimeError(f"Failed to extract archive: {tar_error}") from tar_error else: raise if cleanup_archive: archive_file.unlink(missing_ok=True) print(f"Cleaned up archive: {archive_file}") return extract_to def extract_compose_container_info(container: docker.models.containers.Container, project_name: str) -> QdrantContainer: """Extract container info from a docker-compose container. Args: container: Docker container object from docker-compose project_name: Docker-compose project name Returns: QdrantContainer: Container info object with compose_project set Raises: RuntimeError: If HTTP port mapping cannot be found """ container.reload() port_bindings = container.attrs['NetworkSettings']['Ports'] # Find HTTP port (6333) http_port = None if '6333/tcp' in port_bindings and port_bindings['6333/tcp']: http_port = int(port_bindings['6333/tcp'][0]['HostPort']) else: # Look for any exposed HTTP port for port_key, bindings in port_bindings.items(): if bindings and port_key.endswith('/tcp'): port_num = int(port_key.split('/')[0]) if 6000 <= port_num <= 7000: # Reasonable range for Qdrant http_port = int(bindings[0]['HostPort']) break if not http_port: raise RuntimeError(f"Could not find HTTP port mapping for container {container.name}") # Find gRPC port (6334) - optional grpc_port = None if '6334/tcp' in port_bindings and port_bindings['6334/tcp']: grpc_port = int(port_bindings['6334/tcp'][0]['HostPort']) return QdrantContainer( container=container, host="127.0.0.1", name=container.name, http_port=http_port, grpc_port=grpc_port, compose_project=project_name ) def run_docker_compose(docker_client, qdrant_image, test_data_dir, config): """ Core function to run docker-compose and return container info. Args: docker_client: Docker client instance qdrant_image: Qdrant image to use test_data_dir: Path to test data directory config: Configuration dict with compose_file, wait_for_ready, service_name, and optional `readiness` block forwarded to wait_for_qdrant_ready (keys: scheme, verify, cert, include_grpc). Returns: QdrantDockerCluster: Cluster object containing containers and cleanup function """ wait_for_ready = config.get("wait_for_ready", True) readiness = config.get("readiness") or {} compose_file = config.get("compose_file") if not compose_file: raise ValueError("compose_file parameter is required") # Construct the path to the compose file compose_path = test_data_dir / compose_file if not compose_path.exists(): raise FileNotFoundError(f"Docker compose file not found: {compose_path}") # Generate unique project name to avoid conflicts project_name = f"qdrant-test-{uuid.uuid4().hex[:8]}" service_name = config.get("service_name") # None means return all services # Get docker-compose command compose_cmd = get_docker_compose_command() # Get list of services from compose file services_cmd = compose_cmd + ["-f", str(compose_path), "config", "--services"] result = subprocess.run(services_cmd, capture_output=True, text=True) if result.returncode != 0: raise RuntimeError(f"Failed to get services from compose file: {result.stderr}") services = [s.strip() for s in result.stdout.strip().split('\n') if s.strip()] service_count = len(services) # Start the compose project with custom image override # Set environment variable to override the image in compose services env = dict(os.environ) env["QDRANT_IMAGE"] = qdrant_image compose_up_cmd = compose_cmd + [ "-f", str(compose_path), "-p", project_name, "up", "-d" ] result = subprocess.run(compose_up_cmd, capture_output=True, text=True, env=env) if result.returncode != 0: raise RuntimeError(f"Failed to start docker-compose: {result.stderr}") # Wait for ports to be assigned time.sleep(2) def _readiness_kwargs(info: QdrantContainer) -> Dict[str, Any]: kwargs = {k: v for k, v in readiness.items() if k != "include_grpc"} if readiness.get("include_grpc") and info.grpc_port is not None: kwargs["grpc_port"] = info.grpc_port return {"port": info.http_port, "timeout": 60, **kwargs} if service_count == 1: # Single service compose file - always return single object project_containers = docker_client.containers.list(filters={"name": project_name}) if not project_containers: raise RuntimeError(f"No containers found for project: {project_name}") container = project_containers[0] container_info = extract_compose_container_info(container, project_name) # Wait for this specific container to be ready if wait_for_ready: if not wait_for_qdrant_ready(**_readiness_kwargs(container_info)): raise RuntimeError("Qdrant failed to start within 60 seconds") else: # Multiple services compose file if service_name: # Specific service requested - return single object container_name_prefix = f"{project_name}-{service_name}" containers = docker_client.containers.list(filters={"name": container_name_prefix}) if not containers: raise RuntimeError(f"No container found with prefix: {container_name_prefix}") container = containers[0] container_info = extract_compose_container_info(container, project_name) # Wait for this specific container to be ready if wait_for_ready: if not wait_for_qdrant_ready(**_readiness_kwargs(container_info)): raise RuntimeError("Qdrant failed to start within 60 seconds") else: # No specific service - return array of all container info project_containers = docker_client.containers.list(filters={"name": project_name}) if not project_containers: raise RuntimeError(f"No containers found for project: {project_name}") container_infos = [] for container in project_containers: try: container_info = extract_compose_container_info(container, project_name) container_infos.append(container_info) except Exception as e: print(f"Warning: Could not extract info for container {container.name}: {e}") continue if not container_infos: raise RuntimeError(f"No valid Qdrant containers found in project: {project_name}") # Wait for all containers to be ready if wait_for_ready: for info in container_infos: if not wait_for_qdrant_ready(**_readiness_kwargs(info)): print(f"Warning: Container {info.name} failed to start within 60 seconds") container_info = container_infos # Return the array return QdrantDockerCluster(container_info, project_name, str(compose_path))