test_deploy_models.py 2.8 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162
  1. """Model downloads are stubbed: assert asset scope and failed validation behavior."""
  2. import importlib.util
  3. from pathlib import Path
  4. import sys
  5. import tempfile
  6. from types import SimpleNamespace
  7. import unittest
  8. from unittest.mock import Mock, patch
  9. SPEC = importlib.util.spec_from_file_location('deploy_models', Path(__file__).parents[1] / 'deploy_models.py')
  10. models = importlib.util.module_from_spec(SPEC)
  11. SPEC.loader.exec_module(models)
  12. class DownloadTests(unittest.TestCase):
  13. def setUp(self):
  14. self.temp = tempfile.TemporaryDirectory()
  15. self.addCleanup(self.temp.cleanup)
  16. self.enterContext(patch.object(models, 'ROOT', Path(self.temp.name)))
  17. self.hub = SimpleNamespace(snapshot_download=Mock(), hf_hub_download=Mock())
  18. self.enterContext(patch.dict(sys.modules, huggingface_hub=self.hub))
  19. def test_valid_local_assets_skip_network(self):
  20. with patch.object(models, 'bge_valid', return_value=True), \
  21. patch.object(models, 'tokenizer_valid', return_value=True):
  22. models.download()
  23. self.hub.snapshot_download.assert_not_called()
  24. self.hub.hf_hub_download.assert_not_called()
  25. def test_fresh_server_downloads_only_bge_and_deepseek_tokenizer(self):
  26. with patch.object(models, 'bge_valid', side_effect=[False, True]), \
  27. patch.object(models, 'tokenizer_valid', side_effect=[False, True]):
  28. models.download()
  29. self.assertEqual(self.hub.snapshot_download.call_args.args, ('BAAI/bge-large-zh-v1.5',))
  30. patterns = self.hub.snapshot_download.call_args.kwargs['allow_patterns']
  31. self.assertIn('pytorch_model.bin', patterns)
  32. self.assertNotIn('*.py', patterns)
  33. self.assertEqual(self.hub.hf_hub_download.call_args.args,
  34. ('deepseek-ai/DeepSeek-V3', 'tokenizer.json'))
  35. def test_broken_download_is_not_success(self):
  36. with patch.object(models, 'bge_valid', return_value=False):
  37. with self.assertRaisesRegex(ValueError, 'bge_model_validation_failed'):
  38. models.download()
  39. self.hub.hf_hub_download.assert_not_called()
  40. def test_broken_tokenizer_is_not_success(self):
  41. with patch.object(models, 'bge_valid', return_value=True), \
  42. patch.object(models, 'tokenizer_valid', return_value=False):
  43. with self.assertRaisesRegex(ValueError, 'summary_tokenizer_validation_failed'):
  44. models.download()
  45. def test_active_model_validation_never_repairs_or_downloads(self):
  46. with patch.object(models, 'bge_valid', return_value=False):
  47. with self.assertRaisesRegex(ValueError, 'stop_before_repair'):
  48. models.download(check_only=True)
  49. self.hub.snapshot_download.assert_not_called()
  50. self.hub.hf_hub_download.assert_not_called()
  51. if __name__ == '__main__':
  52. unittest.main()