tests/test_logging_utils.py

from datetime import datetime, timezone, timedelta
from types import SimpleNamespace
from unittest.mock import MagicMock

from huggingface_hub import CommitOperationAdd, CommitOperationDelete
from logging_utils import (
    _build_add_ops, _build_delete_ops, _make_path, _squash_if_needed,
)

NOW = datetime(2026, 5, 17, 12, 0, 0, tzinfo=timezone.utc)
KEEP_COUNT = 10


def _file(days_ago):
    return _make_path(NOW - timedelta(days=days_ago), "abcd1234")


# ── _build_add_ops ────────────────────────────────────────────────────────────

def test_build_add_ops_returns_one_op_per_entry(tmp_path):
    files = [tmp_path / "a.parquet", tmp_path / "b.parquet"]
    for f in files:
        f.write_bytes(b"")
    batch = [("data/a.parquet", str(files[0])), ("data/b.parquet", str(files[1]))]
    ops = _build_add_ops(batch)
    assert len(ops) == 2
    assert all(isinstance(op, CommitOperationAdd) for op in ops)
    assert [op.path_in_repo for op in ops] == ["data/a.parquet", "data/b.parquet"]


def test_build_add_ops_empty_batch():
    assert _build_add_ops([]) == []


# ── _build_delete_ops ─────────────────────────────────────────────────────────

def test_oldest_files_deleted_when_over_limit():
    existing = sorted([_file(i) for i in range(12)])  # 12 files
    ops = _build_delete_ops(existing, n_new=0, max_files=KEEP_COUNT)
    deleted = {op.path_in_repo for op in ops}
    assert len(deleted) == 2
    assert deleted == {_file(11), _file(10)}


def test_files_kept_when_under_limit():
    existing = [_file(i) for i in range(5)]
    assert _build_delete_ops(existing, n_new=0, max_files=KEEP_COUNT) == []


def test_exactly_at_limit_nothing_deleted():
    existing = [_file(i) for i in range(KEEP_COUNT)]
    assert _build_delete_ops(existing, n_new=0, max_files=KEEP_COUNT) == []


def test_one_over_limit_oldest_deleted():
    existing = sorted([_file(i) for i in range(KEEP_COUNT + 1)])
    ops = _build_delete_ops(existing, n_new=0, max_files=KEEP_COUNT)
    assert {op.path_in_repo for op in ops} == {_file(KEEP_COUNT)}


def test_empty_existing_does_nothing():
    assert _build_delete_ops([], n_new=0, max_files=KEEP_COUNT) == []


def test_max_files_zero_skips_pruning():
    existing = [_file(0)]
    assert _build_delete_ops(existing, n_new=0, max_files=0) == []


def test_n_new_counted_toward_total():
    # 8 existing + 4 new = 12 total, need to delete 2
    existing = sorted([_file(i) for i in range(8)])
    ops = _build_delete_ops(existing, n_new=4, max_files=KEEP_COUNT)
    assert len(ops) == 2


def test_oldest_files_deleted_regardless_of_input_order():
    files = [_file(3), _file(11), _file(0), _file(10), _file(1),
             _file(2), _file(4), _file(5), _file(6), _file(7), _file(8)]
    existing = sorted(files)  # caller sorts before passing
    ops = _build_delete_ops(existing, n_new=0, max_files=KEEP_COUNT)
    assert {op.path_in_repo for op in ops} == {_file(11)}


def test_all_ops_are_delete_type():
    existing = sorted([_file(i) for i in range(12)])
    ops = _build_delete_ops(existing, n_new=0, max_files=KEEP_COUNT)
    assert all(isinstance(op, CommitOperationDelete) for op in ops)


# ── _squash_if_needed ─────────────────────────────────────────────────────────

def _squash_api():
    api = MagicMock()
    api.token = "tok"
    return api


def test_squash_runs_when_no_marker(monkeypatch):
    monkeypatch.setattr(
        "logging_utils.hf_hub_download",
        MagicMock(side_effect=FileNotFoundError("no marker")),
    )
    api = _squash_api()
    _squash_if_needed(api, "org/repo")
    api.super_squash_history.assert_called_once()
    api.upload_file.assert_called_once()


def test_squash_skipped_when_marker_is_today(monkeypatch, tmp_path):
    today = datetime.now(timezone.utc).strftime("%Y-%m-%d")
    marker = tmp_path / "last_squash.txt"
    marker.write_text(today)
    monkeypatch.setattr("logging_utils.hf_hub_download", MagicMock(return_value=str(marker)))
    api = _squash_api()
    _squash_if_needed(api, "org/repo")
    api.super_squash_history.assert_not_called()


def test_squash_runs_when_marker_is_yesterday(monkeypatch, tmp_path):
    yesterday = (datetime.now(timezone.utc) - timedelta(days=1)).strftime("%Y-%m-%d")
    marker = tmp_path / "last_squash.txt"
    marker.write_text(yesterday)
    monkeypatch.setattr("logging_utils.hf_hub_download", MagicMock(return_value=str(marker)))
    api = _squash_api()
    _squash_if_needed(api, "org/repo")
    api.super_squash_history.assert_called_once()


def test_squash_error_does_not_raise(monkeypatch):
    monkeypatch.setattr(
        "logging_utils.hf_hub_download",
        MagicMock(side_effect=FileNotFoundError("no marker")),
    )
    api = _squash_api()
    api.super_squash_history.side_effect = RuntimeError("squash failed")
    _squash_if_needed(api, "org/repo")  # must not raise