code/Q-Prefer/tests/test_artifact_resolution.py

from __future__ import annotations

import tempfile
import unittest
from pathlib import Path

from qprefer_reward.scorer import _required_adapter_file


class ArtifactResolutionTest(unittest.TestCase):
    def test_local_required_file_is_returned(self) -> None:
        with tempfile.TemporaryDirectory() as directory:
            expected = Path(directory) / "special_token_embeddings.safetensors"
            expected.write_bytes(b"fixture")
            observed = _required_adapter_file(
                directory,
                "special_token_embeddings.safetensors",
                None,
            )
            self.assertEqual(observed, expected)

    def test_incomplete_local_artifact_fails_fast(self) -> None:
        with (
            tempfile.TemporaryDirectory() as directory,
            self.assertRaisesRegex(FileNotFoundError, "Upload the complete"),
        ):
            _required_adapter_file(
                directory,
                "special_token_embeddings.safetensors",
                None,
            )


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