tests/test_security.py

import tempfile
import unittest
from pathlib import Path

from src.downloads import _safe_filename, _validate_initial_url


class SecurityTests(unittest.TestCase):
    def test_only_https_hf_or_github(self):
        _validate_initial_url("https://huggingface.co/a/b")
        _validate_initial_url("https://github.com/a/b/releases/download/x/model.ckpt")
        with self.assertRaises(ValueError):
            _validate_initial_url("http://github.com/a/b")
        with self.assertRaises(ValueError):
            _validate_initial_url("https://example.com/model.ckpt")

    def test_filename_is_reduced_to_basename(self):
        self.assertEqual(_safe_filename("/a/b/model.ckpt"), "model.ckpt")
        with self.assertRaises(ValueError):
            _safe_filename("...")


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