diff --git a/requirements.txt b/requirements.txt index 3949017..1db1c00 100644 --- a/requirements.txt +++ b/requirements.txt @@ -24,3 +24,4 @@ sentencepiece>=0.1.99 pytest>=7.0.0 pytest-mock>=3.10.0 pytest-cov>=4.0.0 +gguf>=0.6.0 diff --git a/tests/test_service.py b/tests/test_service.py index 3fef3c3..fcb8cf5 100644 --- a/tests/test_service.py +++ b/tests/test_service.py @@ -37,7 +37,7 @@ def test_generate_aibom_basic(self, mock_extractor_cls, mock_score): def test_generate_aibom_purl_encoding(self, mock_extractor_cls, mock_score): # Setup mock_extractor = mock_extractor_cls.return_value - mock_extractor.extract_metadata.return_value = {"name": "test-model", "author": "tester"} + mock_extractor.extract_metadata.return_value = {"name": "test-model", "author": "tester", "commit": "123456"} mock_extractor.extraction_results = {} mock_score.return_value = {"total_score": 50} @@ -48,16 +48,12 @@ def test_generate_aibom_purl_encoding(self, mock_extractor_cls, mock_score): aibom = self.service.generate_aibom(model_id) # Verify PURL encoding (slash should be / now, case preserved) - # Expected: pkg:huggingface/owner/model@12345678 (truncated hash) - - # Check components section (ML model) + # Expected: pkg:huggingface/owner/model@123456 ml_cmp = aibom["components"][0] + + self.assertEqual(ml_cmp["version"], "123456") + self.assertIn("pkg:huggingface/owner/model@123456", ml_cmp["purl"]) self.assertIn("pkg:huggingface/owner/model@123456", ml_cmp["bom-ref"]) - self.assertEqual(ml_cmp["purl"], ml_cmp["bom-ref"]) - - # Check dependency links use matching purls - self.assertIn("pkg:generic/owner/model@123456", aibom["dependencies"][0]["ref"]) - self.assertIn("pkg:huggingface/owner/model@123456", aibom["dependencies"][0]["dependsOn"][0]) @patch("src.models.service.calculate_completeness_score") @patch("src.models.service.EnhancedExtractor")