Support authenticated Hugging Face model and component downloads

test_hub_auth.py

@@ -0,0 +1,44 @@
+import io
+import json
+import tempfile
+import unittest
+from pathlib import Path
+from unittest.mock import patch
+from urllib.request import Request
+from urllib.error import HTTPError
+from hub_auth import HubAuth
+from catalog import Catalog, Redirect, remote
+
+class HubAuthTests(unittest.TestCase):
+ def test_store_replace_remove_and_permissions(self):
+  with tempfile.TemporaryDirectory() as d:
+   a=HubAuth(Path(d)/'key.json');self.assertFalse(a.status()['configured'])
+   a.save('hf_'+'a'*30);self.assertEqual(a.path.stat().st_mode&0o777,0o600)
+   self.assertEqual(a.status(),{'configured':True});a.save('hf_'+'b'*30)
+   self.assertEqual(a.token(),'hf_'+'b'*30);a.save(None);self.assertIsNone(a.token())
+   for bad in ['secret','hf_abc\nAuthorization: bad',{},False]:
+    with self.assertRaises(ValueError):a.save(bad)
+ def test_redirect_never_leaks_to_cdn(self):
+  req=Request('https://huggingface.co/a',headers={'Authorization':'Bearer secret'})
+  handler=Redirect()
+  for url,expected in [('https://cdn-lfs.huggingface.co/b',None),('https://huggingface.co/b','Bearer secret')]:
+   r=handler.redirect_request(req,None,302,'Found',{},url)
+   self.assertEqual(r.get_header('Authorization'),expected)
+  with self.assertRaises(ValueError):handler.redirect_request(req,None,302,'Found',{},'https://evil.test/a')
+ def test_remote_auth_scope_and_safe_errors(self):
+  with patch('catalog.urllib.request.build_opener') as opener:
+   for url,expected in [('https://huggingface.co/a','Bearer secret'),('https://cdn-lfs.huggingface.co/a',None)]:
+    remote(url,'secret');self.assertEqual(opener.return_value.open.call_args.args[0].get_header('Authorization'),expected)
+   for code in [401,403,404,500]:
+    opener.return_value.open.side_effect=HTTPError('https://secret',code,'secret',{},io.BytesIO(b'secret'))
+    with self.assertRaises(ValueError) as err:remote('https://huggingface.co/a','secret')
+    self.assertNotIn('secret',str(err.exception))
+ def test_gated_queue_and_cache(self):
+  with tempfile.TemporaryDirectory() as d:
+   c=Catalog(d);data={'repo':'a/b','revision':'a'*40,'gated':'auto','files':[{'name':'x.gguf','size':1}]}
+   with patch.object(c,'files',return_value=data),patch.object(c,'_next'),patch('catalog.shutil.disk_usage') as disk:
+    disk.return_value.free=100*1024**3
+    with self.assertRaises(ValueError):c.start('a/b','x.gguf','a'*40,'chat')
+    c.cache['a/b']=(0,data);c.hub_credentials('hf_'+'a'*30);self.assertFalse(c.cache)
+    self.assertEqual(c.start('a/b','x.gguf','a'*40,'chat')['state'],'queued')
+    self.assertNotIn('hf_',json.dumps(c.status()))