qdrant/tests/e2e_tests/models.py

"""Data models for e2e tests."""
from typing import Optional, List, Dict, Any, Union
from dataclasses import dataclass, field
import subprocess

import docker.models.containers
import docker.models.networks


class QdrantContainer:
    """Container info object for Qdrant containers."""

    def __init__(self, container: docker.models.containers.Container, host: str, name: str, 
                 http_port: int, grpc_port: Optional[int], compose_project: Optional[str] = None) -> None:
        self.container = container
        self.host = host
        self.name = name
        self.http_port = http_port
        self.grpc_port = grpc_port
        self.compose_project = compose_project

    def __repr__(self) -> str:
        return f"QdrantContainer(name='{self.name}', host='{self.host}', http_port={self.http_port}, grpc_port={self.grpc_port})"


class QdrantCluster:
    """Cluster info object for Qdrant clusters."""

    def __init__(self, leader: QdrantContainer, followers: List[QdrantContainer], 
                 network: docker.models.networks.Network, network_name: str) -> None:
        self.leader = leader
        self.followers = followers
        self.all_nodes = [leader] + followers
        self.network = network
        self.network_name = network_name

    def __repr__(self) -> str:
        return f"QdrantCluster(leader={self.leader.name}, followers={len(self.followers)}, network='{self.network_name}')"


class QdrantDockerCluster:
    """Cluster info object for Docker Compose managed Qdrant clusters."""

    def __init__(self, containers: Union[QdrantContainer, List[QdrantContainer]], 
                 compose_project: str, compose_path: str) -> None:
        """Initialize a Docker Compose cluster.

        Args:
            containers: Single container or list of containers
            compose_project: Docker Compose project name
            compose_path: Path to the docker-compose file
        """
        if isinstance(containers, QdrantContainer):
            self.containers = [containers]
            self.single_container = containers
        else:
            self.containers = containers
            self.single_container = None

        self.compose_project = compose_project
        self.compose_path = compose_path

    def cleanup(self) -> None:
        """Stop and remove the docker-compose deployment including volumes."""
        # Import here to avoid circular dependency
        from .utils import get_docker_compose_command

        compose_cmd = get_docker_compose_command()
        compose_down_cmd = compose_cmd + [
            "-f", str(self.compose_path),
            "-p", self.compose_project,
            "down", "-v"  # Also remove volumes
        ]

        result = subprocess.run(compose_down_cmd, capture_output=True, text=True)
        if result.returncode != 0:
            print(f"Warning: Failed to stop docker-compose: {result.stderr}")
        else:
            print(f"Cleaned up compose project: {self.compose_project}")

    def get_container(self, index: int = 0) -> QdrantContainer:
        """Get a specific container by index."""
        if index >= len(self.containers):
            raise IndexError(f"Container index {index} out of range (have {len(self.containers)} containers)")
        return self.containers[index]

    def get_all_containers(self) -> List[QdrantContainer]:
        """Get all containers in the cluster."""
        return self.containers

    def __len__(self) -> int:
        """Return the number of containers in the cluster."""
        return len(self.containers)

    def __getitem__(self, index: int) -> QdrantContainer:
        """Support indexing to get containers."""
        return self.containers[index]

    def __iter__(self):
        """Support iteration over containers."""
        return iter(self.containers)

    def __repr__(self) -> str:
        container_names = ", ".join(c.name for c in self.containers)
        return f"QdrantDockerCluster(project='{self.compose_project}', containers=[{container_names}])"


@dataclass
class QdrantContainerConfig:
    """Configuration for creating Qdrant containers with proper typing and defaults."""

    # Container identification
    name: Optional[str] = None

    # Resource limits
    mem_limit: Optional[str] = None
    cpu_limit: Optional[str] = None

    # Network configuration
    network: Optional[str] = None
    network_mode: Optional[str] = None

    # Storage configuration
    volumes: Optional[Dict[str, Dict[str, str]]] = None
    mounts: Optional[List[Any]] = None

    # Environment and command
    environment: Optional[Dict[str, str]] = None
    command: Optional[Union[str, List[str]]] = None

    # Container behavior
    remove: bool = True
    detach: bool = True

    # Qdrant-specific settings
    exit_on_error: bool = True

    # Additional Docker parameters
    additional_params: Optional[Dict[str, Any]] = field(default_factory=dict)

    def to_docker_config(self, qdrant_image: str) -> Dict[str, Any]:
        """Convert this config to Docker container.run() parameters."""
        config = {
            "image": qdrant_image,
            "detach": self.detach,
            "remove": self.remove,
        }

        # Add port bindings unless host network mode
        if self.network_mode != 'host':
            config["ports"] = {'6333/tcp': ('127.0.0.1', None), '6334/tcp': ('127.0.0.1', None)}

        # Add optional parameters
        if self.name:
            config["name"] = self.name
        if self.mem_limit:
            config["mem_limit"] = self.mem_limit
        if self.cpu_limit:
            config["cpu_limit"] = self.cpu_limit
        if self.network:
            config["network"] = self.network
        if self.network_mode:
            config["network_mode"] = self.network_mode
        if self.volumes:
            config["volumes"] = self.volumes
        if self.mounts:
            config["mounts"] = self.mounts
        if self.environment:
            config["environment"] = self.environment
        if self.command:
            config["command"] = self.command

        # Add any additional parameters
        if self.additional_params:
            config.update(self.additional_params)

        return config