tests/conftest.py

"""Shared test fixtures."""

from __future__ import annotations

import tempfile
from pathlib import Path
from typing import TYPE_CHECKING

import nibabel as nib
import numpy as np
import pytest

from stroke_deepisles_demo.core.types import CaseFiles

if TYPE_CHECKING:
    from collections.abc import Generator


@pytest.fixture
def temp_dir() -> Generator[Path, None, None]:
    """Create a temporary directory for test outputs."""
    with tempfile.TemporaryDirectory() as td:
        yield Path(td)


@pytest.fixture
def synthetic_nifti_3d(temp_dir: Path) -> Path:
    """Create a minimal synthetic 3D NIfTI file."""
    data = np.random.rand(10, 10, 10).astype(np.float32)
    img = nib.Nifti1Image(data, affine=np.eye(4))  # type: ignore
    path = temp_dir / "synthetic.nii.gz"
    nib.save(img, path)  # type: ignore
    return path


@pytest.fixture
def synthetic_case_files(temp_dir: Path) -> CaseFiles:
    """Create a complete set of synthetic case files."""
    # Create DWI
    dwi_data = np.random.rand(64, 64, 30).astype(np.float32)
    dwi_img = nib.Nifti1Image(dwi_data, affine=np.eye(4))  # type: ignore
    dwi_path = temp_dir / "dwi.nii.gz"
    nib.save(dwi_img, dwi_path)  # type: ignore

    # Create ADC
    adc_data = np.random.rand(64, 64, 30).astype(np.float32) * 2000
    adc_img = nib.Nifti1Image(adc_data, affine=np.eye(4))  # type: ignore
    adc_path = temp_dir / "adc.nii.gz"
    nib.save(adc_img, adc_path)  # type: ignore

    # Create mask
    mask_data = (np.random.rand(64, 64, 30) > 0.9).astype(np.uint8)
    mask_img = nib.Nifti1Image(mask_data, affine=np.eye(4))  # type: ignore
    mask_path = temp_dir / "mask.nii.gz"
    nib.save(mask_img, mask_path)  # type: ignore

    return CaseFiles(
        dwi=dwi_path,
        adc=adc_path,
        ground_truth=mask_path,
    )


@pytest.fixture
def synthetic_isles_dir(temp_dir: Path) -> Path:
    """
    Create synthetic ISLES24-like directory structure.

    Structure:
        temp_dir/
        ├── Images-DWI/
        │   ├── sub-stroke0001_ses-02_dwi.nii.gz
        │   └── sub-stroke0002_ses-02_dwi.nii.gz
        ├── Images-ADC/
        │   ├── sub-stroke0001_ses-02_adc.nii.gz
        │   └── sub-stroke0002_ses-02_adc.nii.gz
        └── Masks/
            ├── sub-stroke0001_ses-02_lesion-msk.nii.gz
            └── sub-stroke0002_ses-02_lesion-msk.nii.gz
    """
    dwi_dir = temp_dir / "Images-DWI"
    adc_dir = temp_dir / "Images-ADC"
    mask_dir = temp_dir / "Masks"

    dwi_dir.mkdir()
    adc_dir.mkdir()
    mask_dir.mkdir()

    for subject_num in [1, 2]:
        subject_id = f"sub-stroke{subject_num:04d}"

        # Create DWI
        dwi_data = np.random.rand(10, 10, 5).astype(np.float32)
        dwi_img = nib.Nifti1Image(dwi_data, affine=np.eye(4))  # type: ignore
        nib.save(dwi_img, dwi_dir / f"{subject_id}_ses-02_dwi.nii.gz")  # type: ignore

        # Create ADC
        adc_data = np.random.rand(10, 10, 5).astype(np.float32) * 2000
        adc_img = nib.Nifti1Image(adc_data, affine=np.eye(4))  # type: ignore
        nib.save(adc_img, adc_dir / f"{subject_id}_ses-02_adc.nii.gz")  # type: ignore

        # Create Mask
        mask_data = (np.random.rand(10, 10, 5) > 0.9).astype(np.uint8)
        mask_img = nib.Nifti1Image(mask_data, affine=np.eye(4))  # type: ignore
        nib.save(mask_img, mask_dir / f"{subject_id}_ses-02_lesion-msk.nii.gz")  # type: ignore

    return temp_dir