tests/test_openclip_checkpoint_path.py

from __future__ import annotations

import os
import sys
import tempfile
import types
import unittest
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import Mock, patch

from satquery_agent.satquery_vision_encoder import (
    OPENCLIP_IMAGE_INTERPOLATION,
    OPENCLIP_IMAGE_MEAN,
    OPENCLIP_IMAGE_RESIZE_MODE,
    OPENCLIP_IMAGE_STD,
    SatQueryVisionEncoder,
    resolve_openclip_checkpoint,
)
from satquery_agent.sve_artifacts import (
    EXPECTED_EMBEDDING_DIMENSION,
    EXPECTED_OPENCLIP_FILENAME,
    EXPECTED_OPENCLIP_MODEL,
    EXPECTED_OPENCLIP_REPOSITORY,
    EXPECTED_OPENCLIP_REVISION,
    EXPECTED_PRETRAINED,
)


class _FakeOpenCLIPModel:
    def state_dict(self):
        return {}

    def load_state_dict(self, state, strict=False):
        return SimpleNamespace(unexpected_keys=[], missing_keys=[])

    def eval(self):
        return self

    def parameters(self):
        return []

    def to(self, *args, **kwargs):
        return self


class OpenCLIPCheckpointPathTests(unittest.TestCase):
    def test_extensionless_hf_blob_is_passed_to_openclip_as_safetensors(self):
        import torch

        with tempfile.TemporaryDirectory() as temporary_directory:
            root = Path(temporary_directory)
            blob = root / "hub" / "models--laion--clip" / "blobs" / "7d129ed"
            blob.parent.mkdir(parents=True)
            blob.write_bytes(b"safetensors checkpoint bytes")
            snapshot = (
                root
                / "hub"
                / "models--laion--clip"
                / "snapshots"
                / EXPECTED_OPENCLIP_REVISION
                / EXPECTED_OPENCLIP_FILENAME
            )
            snapshot.parent.mkdir(parents=True)
            snapshot.symlink_to(blob)
            self.assertTrue(snapshot.resolve().samefile(blob))
            self.assertEqual(snapshot.resolve().suffix, "")

            download = Mock(return_value=str(snapshot))
            received = {}
            fake_open_clip = types.ModuleType("open_clip")

            def create_model_and_transforms(model_name, **kwargs):
                received["model_name"] = model_name
                received.update(kwargs)
                return _FakeOpenCLIPModel(), object(), object()

            fake_open_clip.create_model_and_transforms = create_model_and_transforms
            fake_open_clip.get_tokenizer = lambda model_name: object()

            adapter_payload = {
                "adapter_state_dict": {},
                "model_name": "SatQuery Vision Encoder v1",
                "backbone": EXPECTED_OPENCLIP_MODEL,
                "pretrained": EXPECTED_PRETRAINED,
                "embedding_dimension": EXPECTED_EMBEDDING_DIMENSION,
                "unfreeze_visual_blocks": 2,
                "adaptation_dataset": "BigEarthNet.txt",
                "image_dataset": "BigEarthNet v2 Lithuania Summer",
                "train_samples": 4008,
                "validation_samples": 2291,
                "test_samples": 2053,
                "epoch": 5,
            }
            artifacts = SimpleNamespace(adapter_path=root / "adapter.pt")

            with (
                patch.dict(os.environ, {"SATQUERY_MODEL_CACHE": str(root / "cache")}),
                patch.dict(sys.modules, {"open_clip": fake_open_clip}),
                patch("huggingface_hub.hf_hub_download", download),
                patch.object(torch, "load", return_value=adapter_payload) as load_adapter,
            ):
                encoder = SatQueryVisionEncoder(artifacts, "cpu")
                encoder.load()

                checkpoint = Path(received["pretrained"])
                self.assertEqual(checkpoint.name, EXPECTED_OPENCLIP_FILENAME)
                self.assertEqual(checkpoint.suffix, ".safetensors")
                self.assertEqual(checkpoint.read_bytes(), blob.read_bytes())
                self.assertEqual(received["model_name"], EXPECTED_OPENCLIP_MODEL)
                self.assertEqual(received["device"], "cpu")
                self.assertEqual(received["image_mean"], OPENCLIP_IMAGE_MEAN)
                self.assertEqual(received["image_std"], OPENCLIP_IMAGE_STD)
                self.assertEqual(received["image_interpolation"], OPENCLIP_IMAGE_INTERPOLATION)
                self.assertEqual(received["image_resize_mode"], OPENCLIP_IMAGE_RESIZE_MODE)
                download.assert_called_once_with(
                    repo_id=EXPECTED_OPENCLIP_REPOSITORY,
                    filename=EXPECTED_OPENCLIP_FILENAME,
                    revision=EXPECTED_OPENCLIP_REVISION,
                    cache_dir=str((root / "cache").resolve()),
                    token=False,
                )
                load_adapter.assert_called_once_with(
                    artifacts.adapter_path,
                    map_location="cpu",
                    weights_only=True,
                )

                downloader_that_must_not_run = Mock(side_effect=AssertionError("unexpected download"))
                self.assertEqual(
                    resolve_openclip_checkpoint(downloader=downloader_that_must_not_run),
                    checkpoint,
                )
                downloader_that_must_not_run.assert_not_called()


if __name__ == "__main__":
    unittest.main()