File size: 824 Bytes
81ba775 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 | 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()
|