visual-semiotic-ai-analysis/tests/integration/conftest.py

"""Integration test configuration."""

import logging
import os
import shutil
import tempfile
from pathlib import Path
import json
from collections.abc import Generator
from uuid import UUID
from io import BytesIO
import pandas as pd
from PIL import Image
import pytest
import structlog
from minio import Minio, S3Error
from testcontainers.minio import MinioContainer
from python_repositories.adapters import MinioAdapter
from data_store.repositories import (
    ImageRepository,
    AnnotationRepository,
    JobRepository,
)
from data_store.dto import (
    Annotation,
    DistanceEnum,
    AngleEnum,
    PointOfViewEnum,
    ContactEnum,
    Job,
)

# Last Docker image with a useful WebUI
MINIO_DOCKER_IMAGE = "minio/minio:RELEASE.2023-03-20T20-16-18Z"


@pytest.fixture(scope="session", autouse=True)
def configure_logging() -> None:
    """Configure logging for the test session."""
    # Configure structlog
    structlog.configure(
        processors=[
            structlog.stdlib.filter_by_level,
            structlog.stdlib.add_logger_name,
            structlog.stdlib.add_log_level,
            structlog.processors.TimeStamper(fmt="iso"),
            structlog.processors.JSONRenderer(),
        ],
        logger_factory=structlog.stdlib.LoggerFactory(),
        wrapper_class=structlog.stdlib.BoundLogger,
        cache_logger_on_first_use=True,
    )
    # Set up basic logging configuration
    logging.basicConfig(level=logging.ERROR)


@pytest.fixture(scope="session")
def minio_container() -> Generator[dict[str, str]]:
    """Set up MinIO container and provide credentials for testing."""
    # Set variables
    access_key = "minioadmin"
    secret_key = "minioadmin"
    bucket_name = "test-bucket"
    # Start container
    container = MinioContainer(
        image=MINIO_DOCKER_IMAGE,
        access_key=access_key,
        secret_key=secret_key,
    ).with_bind_ports(  # expose webUI on 9001
        container=9001,
        host=9001,
    )
    container.start()
    # Build environment variables dictionary
    minio_host = container.get_container_host_ip()
    minio_port = container.get_exposed_port(9000)
    minio_endpoint = f"{minio_host}:{minio_port}"
    env_vars = {
        "MINIO_ENDPOINT": minio_endpoint,
        "MINIO_ACCESS_KEY": access_key,
        "MINIO_SECRET_KEY": secret_key,
        "MINIO_BUCKET": bucket_name,
    }

    yield env_vars

    # Stop container
    container.stop()


@pytest.fixture(scope="session", autouse=True)
def set_environment_variables(
    minio_container: dict[str, str],
) -> Generator[dict[str, str]]:
    """Set environment variables needed for tests."""
    # Build environment variables dictionary
    env_vars = {}
    env_vars.update(minio_container)
    # Set environment variables
    for key, value in env_vars.items():
        os.environ[key] = value

    yield env_vars

    # Cleanup
    for key in env_vars:
        _ = os.environ.pop(key, default=None)


@pytest.fixture(scope="session")
def job_id_list() -> Generator[list[UUID]]:
    """Provide a list of job IDs for testing."""
    yield [
        UUID("d358cbba-75be-4c38-becd-7abf338ca941"),
        UUID("56d20d25-b733-4ed3-880c-732d7257dcaa"),
        UUID("8019292f-8353-43d7-a189-076c6630de7b"),
    ]


@pytest.fixture(scope="session")
def csv_data_file(
    job_id_list: list[UUID],
) -> Generator[Path]:
    """Create a CSV data file for testing."""
    # Create csv content
    csv_data = pd.DataFrame(
        {
            "job_id": job_id_list,
            "distance": [
                "close",
                "long",
                "wrong",
            ],
            "angle": ["high", "low", "wrong"],
            "point_of_view": ["frontal", "oblique", "wrong"],
            "contact": ["offer", "demand", "wrong"],
        }
    )
    # Create data directory and file path
    data_dir = Path(tempfile.mkdtemp())
    data_path = data_dir / "annotations.csv"
    # Write CSV content to file
    csv_data.to_csv(data_path, header=True, index=False)

    yield data_path

    # Cleanup
    shutil.rmtree(data_dir)


@pytest.fixture(scope="session")
def raw_minio_client(
    set_environment_variables: dict[str, str],
) -> Generator[Minio]:
    """Provide a MinIO client for testing."""
    # Create MinIO client
    client = Minio(
        endpoint=set_environment_variables["MINIO_ENDPOINT"],
        access_key=set_environment_variables["MINIO_ACCESS_KEY"],
        secret_key=set_environment_variables["MINIO_SECRET_KEY"],
        secure=False,
    )
    # Create bucket if it doesn't exist
    bucket_name = set_environment_variables["MINIO_BUCKET"]
    try:
        if not client.bucket_exists(bucket_name):
            client.make_bucket(bucket_name)
    except S3Error as e:
        raise RuntimeError(f"Failed to create or access bucket: {e}")

    yield client

    # No explicit cleanup needed for MinIO client


@pytest.fixture(scope="session")
def image_repository() -> Generator[ImageRepository]:
    """Provide an ImageRepository instance for testing."""
    # Instantiate and connect repository
    repo = ImageRepository()
    repo.connect()

    yield repo

    # Cleanup
    repo.disconnect()


@pytest.fixture(scope="session")
def image() -> Generator[Image.Image]:
    """Provide a sample image for testing."""
    img = Image.new("RGB", (10, 10))
    yield img


@pytest.fixture(scope="function")
def image_in_minio(
    set_environment_variables: dict[str, str],
    raw_minio_client: Minio,
    image: Image.Image,
    job_id_list: list[UUID],
) -> Generator[tuple[UUID, Image.Image]]:
    """Upload a sample image to MinIO and provide its ID for testing."""
    # Prepare object name and buffer
    bucket_name = set_environment_variables["MINIO_BUCKET"]
    image_id = job_id_list[0]
    object_name = ImageRepository._object_name(image_id)
    buffer = BytesIO()
    image.save(buffer, ImageRepository.image_format)
    num_bytes = buffer.getbuffer().nbytes
    buffer.seek(0)
    # Upload image to MinIO
    raw_minio_client.put_object(
        bucket_name=bucket_name,
        object_name=object_name,
        data=buffer,
        length=num_bytes,
        part_size=MinioAdapter.chunk_size,
    )
    buffer.close()

    yield image_id, image

    # Cleanup: remove image from MinIO
    try:
        raw_minio_client.remove_object(
            bucket_name=bucket_name,
            object_name=object_name,
        )
    except S3Error as e:
        raise RuntimeError(f"Failed to remove object from bucket: {e}")


@pytest.fixture(scope="session")
def annotation_repository() -> Generator[AnnotationRepository]:
    """Provide an AnnotationRepository instance for testing."""
    # Instantiate and connect repository
    repo = AnnotationRepository()
    repo.connect()

    yield repo

    # Cleanup
    repo.disconnect()


@pytest.fixture(scope="session")
def annotation() -> Generator[Annotation]:
    """Provide a sample annotation for testing."""
    annotation = Annotation(
        distance=DistanceEnum.CLOSE,
        angle=AngleEnum.HIGH,
        point_of_view=PointOfViewEnum.FRONTAL,
        contact=ContactEnum.OFFER,
    )
    yield annotation


@pytest.fixture(scope="function")
def annotation_in_minio(
    set_environment_variables: dict[str, str],
    raw_minio_client: Minio,
    annotation: Annotation,
    job_id_list: list[UUID],
) -> Generator[tuple[UUID, Annotation]]:
    """Upload a sample annotation to MinIO and provide its ID for testing."""
    # Prepare object name and buffer
    bucket_name = set_environment_variables["MINIO_BUCKET"]
    job_id = job_id_list[0]
    object_name = AnnotationRepository._object_name(job_id)
    data_json = annotation.model_dump(mode="json")
    data_bytes = json.dumps(data_json).encode(AnnotationRepository.encoding)
    buffer = BytesIO(data_bytes)
    num_bytes = buffer.getbuffer().nbytes
    buffer.seek(0)
    # Upload annotation to MinIO
    raw_minio_client.put_object(
        bucket_name=bucket_name,
        object_name=object_name,
        data=buffer,
        length=num_bytes,
        part_size=MinioAdapter.chunk_size,
    )
    buffer.close()

    yield job_id, annotation

    # Cleanup: remove annotation from MinIO
    try:
        raw_minio_client.remove_object(
            bucket_name=bucket_name,
            object_name=object_name,
        )
    except S3Error as e:
        raise RuntimeError(f"Failed to remove object from bucket: {e}")


@pytest.fixture(scope="session")
def job_repository() -> Generator[JobRepository]:
    """Provide a JobRepository instance for testing."""
    # Instantiate and connect repository
    repo = JobRepository()
    repo.connect()

    yield repo

    # Cleanup
    repo.disconnect()


@pytest.fixture(scope="session")
def job(
    job_id_list: list[UUID],
    image: Image.Image,
    annotation: Annotation,
) -> Generator[Job]:
    """Provide a sample job for testing."""
    job = Job(
        job_id=job_id_list[0],
        image=image,
        annotation=annotation,
    )
    yield job


@pytest.fixture(scope="function")
def job_in_minio(
    job: Job,
    image_in_minio: tuple[UUID, Image.Image],
    annotation_in_minio: tuple[UUID, Annotation],
) -> Generator[Job]:
    """Provide a sample job with its components already in MinIO for testing."""
    # A job consists of an image and an annotation and both must be in MinIO
    # before the job can be considered to be in MinIO.
    # Ensure the job_id matches those of the image and annotation
    image_id, _ = image_in_minio
    annotation_id, _ = annotation_in_minio
    assert job.job_id == annotation_id
    assert job.job_id == image_id
    yield job