code/tests/conftest.py

"""Shared pytest fixtures."""

from __future__ import annotations

import mfflora._compat  # noqa: F401

import os
from pathlib import Path

import pytest
import torch


@pytest.fixture(autouse=True)
def _seed_everything():
    torch.manual_seed(0)


@pytest.fixture
def factor_library_root() -> Path:
    return Path(
        os.environ.get(
            "MFFLORA_FACTOR_ROOT",
            "external/kronlingua/factors/e1_step2_llama31_base_12lang_gate",
        )
    )


@pytest.fixture
def has_factor_library(factor_library_root: Path) -> bool:
    return (factor_library_root / "en" / "model__layers__0__mlp__gate_proj.safetensors").exists()


def _random_psd(n: int, *, seed: int = 0, scale: float = 1.0) -> torch.Tensor:
    g = torch.Generator().manual_seed(seed)
    L = torch.randn(n, n, generator=g) * scale
    return L @ L.T + 1e-3 * torch.eye(n)


@pytest.fixture
def random_psd_factory():
    return _random_psd