social-data-attribution-w2/tests/test_query_data_scores.py

from __future__ import annotations

import json
import sys
from pathlib import Path

import numpy as np
import pytest


REPO_ROOT = Path(__file__).resolve().parents[1]
QUERY_DATA_DIR = REPO_ROOT / "scripts" / "query_data"

if str(QUERY_DATA_DIR) not in sys.path:
    sys.path.insert(0, str(QUERY_DATA_DIR))

import merge_scores  # noqa: E402


def _write_score_dir(path: Path, values: list[list[float]]) -> None:
    num_scores = len(values)
    num_items = len(values[0])
    dtype_json = {
        "names": [
            item for i in range(num_scores) for item in (f"score_{i}", f"written_{i}")
        ],
        "formats": [item for _ in range(num_scores) for item in ("float32", "bool")],
        "offsets": [item for i in range(num_scores) for item in (i * 8, i * 8 + 4)],
        "itemsize": num_scores * 8,
    }
    dtype = np.dtype(dtype_json)
    mmap = np.memmap(path / "scores.bin", dtype=dtype, mode="w+", shape=(num_items,))
    for index, column in enumerate(values):
        mmap[f"score_{index}"] = np.array(column, dtype=np.float32)
        mmap[f"written_{index}"] = True
    mmap.flush()

    info = {
        "num_items": num_items,
        "num_scores": num_scores,
        "dtype": dtype_json,
    }
    (path / "info.json").write_text(json.dumps(info), encoding="utf-8")


def test_merge_scores_load_fp32_scores(tmp_path: Path) -> None:
    _write_score_dir(tmp_path, [[1.5, 3.0, -2.0]])

    merge_loaded, num_items = merge_scores.load_scores(tmp_path)

    assert num_items == 3
    assert merge_loaded.tolist() == pytest.approx([1.5, 3.0, -2.0])


def test_merge_scores_reject_multi_score_outputs(tmp_path: Path) -> None:
    _write_score_dir(tmp_path, [[1.0, 2.0], [3.0, 4.0]])

    with pytest.raises(ValueError, match="expected 1"):
        merge_scores.load_scores(tmp_path)