OBLITERATUS/tests/test_persistence_contracts.py

"""Boundary contracts for deterministic and transactional checkpoint state."""

from __future__ import annotations

import json
import os
from pathlib import Path
from unittest.mock import MagicMock

import pytest
import torch

import obliteratus.persistence_contracts as persistence


pytestmark = pytest.mark.cpu

@pytest.mark.parametrize(
    ("state_dict", "expected"),
    [
        ({}, 0),
        ({"float": torch.ones(3, dtype=torch.float32)}, 12),
        (
            {
                "half": torch.ones(5, dtype=torch.float16),
                "index": torch.ones(2, dtype=torch.int64),
            },
            26,
        ),
    ],
)
def test_state_dict_size_bytes_is_exact(state_dict, expected):
    assert persistence.state_dict_size_bytes(state_dict) == expected


@pytest.mark.parametrize(
    ("payload", "required"),
    [(0, 0), (1, 2), (9, 10), (10, 11), (101, 112)],
)
def test_required_checkpoint_bytes_rounds_headroom_up(payload, required):
    assert persistence.required_checkpoint_bytes(payload) == required


def test_required_checkpoint_bytes_rejects_negative_payload():
    with pytest.raises(ValueError) as exc_info:
        persistence.required_checkpoint_bytes(-1)

    assert str(exc_info.value) == "checkpoint payload size cannot be negative"


def test_checkpoint_capacity_accepts_exact_boundary_and_rejects_one_byte_less():
    persistence.ensure_checkpoint_capacity(free_bytes=11, payload_bytes=10)

    with pytest.raises(
        OSError,
        match=r"Insufficient disk space: 0\.0 GB free, need ~0\.0 GB",
    ):
        persistence.ensure_checkpoint_capacity(free_bytes=10, payload_bytes=10)


def test_checkpoint_capacity_error_reports_decimal_gigabytes_exactly():
    payload_bytes = 1_000_000_000_000_000_000

    with pytest.raises(OSError) as exc_info:
        persistence.ensure_checkpoint_capacity(
            free_bytes=payload_bytes,
            payload_bytes=payload_bytes,
        )

    assert str(exc_info.value) == (
        "Insufficient disk space: 1000000000.0 GB free, "
        "need ~1000000000.0 GB. "
        "Try a different --output-dir on a larger filesystem."
    )


def test_metadata_serialization_is_stable_strict_json():
    first = persistence.serialize_checkpoint_metadata({"z": 1, "a": [True, None]})
    second = persistence.serialize_checkpoint_metadata({"a": [True, None], "z": 1})

    assert first == second
    assert json.loads(first) == {"a": [True, None], "z": 1}
    assert first == (
        '{\n  "a": [\n    true,\n    null\n  ],\n  "z": 1\n}'
    )


def test_metadata_serializer_passes_strict_format_options(monkeypatch):
    real_dumps = json.dumps
    observed = {}

    def record_dumps(metadata, **kwargs):
        observed.update(kwargs)
        return real_dumps(metadata, **kwargs)

    monkeypatch.setattr(persistence.json, "dumps", record_dumps)

    assert json.loads(persistence.serialize_checkpoint_metadata({"schema": 1})) == {
        "schema": 1,
    }
    assert observed == {"indent": 2, "sort_keys": True, "allow_nan": False}


@pytest.mark.parametrize("value", [object(), float("nan"), float("inf")])
def test_metadata_serialization_rejects_nonportable_values(value):
    with pytest.raises((TypeError, ValueError)):
        persistence.serialize_checkpoint_metadata({"invalid": value})


def test_remove_checkpoint_path_handles_file_directory_symlink_and_missing(tmp_path):
    file_path = tmp_path / "file"
    file_path.write_text("data", encoding="utf-8")
    directory = tmp_path / "directory"
    directory.mkdir()
    (directory / "nested").write_text("data", encoding="utf-8")
    target = tmp_path / "target"
    target.mkdir()
    (target / "sentinel").write_text("preserve", encoding="utf-8")
    symlink = tmp_path / "link"
    symlink.symlink_to(target, target_is_directory=True)

    for path in (file_path, directory, symlink, tmp_path / "missing"):
        persistence._remove_checkpoint_path(path)

    assert not file_path.exists()
    assert not directory.exists()
    assert not symlink.exists()
    assert (target / "sentinel").read_text(encoding="utf-8") == "preserve"


def test_remove_checkpoint_path_requests_race_safe_file_unlink():
    path = MagicMock()
    path.is_symlink.return_value = True

    persistence._remove_checkpoint_path(path)

    path.unlink.assert_called_once_with(missing_ok=True)
    path.is_file.assert_not_called()


def test_remove_checkpoint_path_requests_best_effort_directory_cleanup(monkeypatch):
    path = MagicMock()
    path.is_symlink.return_value = False
    path.is_file.return_value = False
    path.exists.return_value = True
    remove_tree = MagicMock()
    monkeypatch.setattr(persistence.shutil, "rmtree", remove_tree)

    persistence._remove_checkpoint_path(path)

    remove_tree.assert_called_once_with(path, ignore_errors=True)


def test_atomic_checkpoint_creates_parent_and_promotes_new_destination(tmp_path):
    destination = tmp_path / "nested" / "checkpoint"

    with persistence.atomic_checkpoint_directory(destination) as staging:
        assert staging.parent == destination.parent
        assert staging.name.startswith(".checkpoint.staging-")
        (staging / "model.bin").write_bytes(b"complete")

    assert (destination / "model.bin").read_bytes() == b"complete"
    assert list(destination.parent.glob(".checkpoint.*-*")) == []


def test_atomic_checkpoint_cleans_staging_when_new_destination_promotion_fails(
    tmp_path,
    monkeypatch,
):
    destination = tmp_path / "checkpoint"

    def fail_promotion(source, target):
        raise OSError(f"cannot promote {Path(source).name} to {Path(target).name}")

    monkeypatch.setattr(persistence.os, "replace", fail_promotion)
    with pytest.raises(OSError, match="cannot promote"):
        with persistence.atomic_checkpoint_directory(destination) as staging:
            (staging / "model.bin").write_bytes(b"partial")

    assert not destination.exists()
    assert list(tmp_path.glob(".checkpoint.staging-*")) == []


def test_atomic_checkpoint_replaces_symlink_without_touching_target(tmp_path):
    target = tmp_path / "target"
    target.mkdir()
    sentinel = target / "sentinel"
    sentinel.write_text("preserve", encoding="utf-8")
    destination = tmp_path / "checkpoint"
    destination.symlink_to(target, target_is_directory=True)

    with persistence.atomic_checkpoint_directory(destination) as staging:
        (staging / "model.bin").write_bytes(b"replacement")

    assert not destination.is_symlink()
    assert (destination / "model.bin").read_bytes() == b"replacement"
    assert sentinel.read_text(encoding="utf-8") == "preserve"


def test_atomic_checkpoint_uses_os_replace_not_copy(tmp_path, monkeypatch):
    destination = tmp_path / "checkpoint"
    calls = []
    real_replace = os.replace

    def record_replace(source, target):
        calls.append((Path(source), Path(target)))
        return real_replace(source, target)

    monkeypatch.setattr(persistence.os, "replace", record_replace)
    with persistence.atomic_checkpoint_directory(destination) as staging:
        (staging / "model.bin").write_bytes(b"saved")

    assert len(calls) == 1
    assert calls[0][1] == destination
    assert ".staging-" in calls[0][0].name