tests/test_checkpointing.py

import pytest
import torch

import src.model.encoder.depth.depthpro.depthpro_wrapper as depthpro_wrapper
from src.demo import infer_single_image
from src.demo.infer_single_image import (
    _extract_state_dict,
    _resolve_checkpoint_path,
    load_demo_encoder,
)


def test_extract_state_dict_uses_safe_loading(tmp_path, monkeypatch) -> None:
    checkpoint = tmp_path / "model.ckpt"
    expected = {"encoder.weight": torch.tensor([1.0])}
    torch.save(expected, checkpoint)

    original_load = torch.load

    def tracked_load(*args, **kwargs):
        assert kwargs["weights_only"] is True
        return original_load(*args, **kwargs)

    monkeypatch.setattr(torch, "load", tracked_load)

    actual = _extract_state_dict(checkpoint)

    assert actual.keys() == expected.keys()
    assert torch.equal(actual["encoder.weight"], expected["encoder.weight"])


def test_resolve_checkpoint_path_accepts_local_file(tmp_path) -> None:
    checkpoint = tmp_path / "model.ckpt"
    checkpoint.touch()

    assert _resolve_checkpoint_path(str(checkpoint)) == checkpoint


def test_resolve_checkpoint_path_rejects_missing_file(tmp_path) -> None:
    checkpoint = tmp_path / "missing.ckpt"

    with pytest.raises(FileNotFoundError, match="Checkpoint not found"):
        _resolve_checkpoint_path(str(checkpoint))


def test_load_demo_encoder_ignores_decoder_weights(tmp_path, monkeypatch) -> None:
    checkpoint = tmp_path / "model.ckpt"
    torch.save(
        {
            "encoder.weight": torch.tensor([[3.0]]),
            "decoder.unused": torch.tensor([7.0]),
        },
        checkpoint,
    )
    encoder = torch.nn.Linear(1, 1, bias=False)
    monkeypatch.setattr(infer_single_image, "get_encoder", lambda _: encoder)

    loaded = load_demo_encoder(
        cfg=type("Cfg", (), {"model": type("Model", (), {"encoder": object()})()})(),
        checkpoint_path=checkpoint,
        device=torch.device("cpu"),
    )

    assert loaded is encoder
    assert not loaded.training
    assert torch.equal(loaded.weight, torch.tensor([[3.0]]))


def test_depthpro_does_not_require_external_checkpoint(monkeypatch) -> None:
    captured = {}

    def create_without_weights(config, **_kwargs):
        captured["checkpoint_uri"] = config.checkpoint_uri
        return torch.nn.Identity(), None

    monkeypatch.setattr(
        depthpro_wrapper,
        "create_model_and_transforms",
        create_without_weights,
    )

    depthpro_wrapper.DepthPro()

    assert captured["checkpoint_uri"] is None