files/tests/test_utils.py

@@ -1,7 +1,10 @@
 from pathlib import Path
 from unittest.mock import patch, ANY
 import dataclasses
+import io
+import os
 import tempfile
+import zipfile
 
 from django.test import TestCase
 
@@ -11,8 +14,10 @@ from files.utils import (
     find_exact_path,
     find_path_by_name,
     get_thumbnail_upload_to,
+    get_wheels_from_manifest,
     make_thumbnails,
     validate_file_list,
+    validate_wheels,
 )
 
 # Reusing test files from the extensions app
@@ -290,3 +295,62 @@ class UtilsTest(TestCase):
                     validate_file_list(test.toml_content, test.manifest_filepath, test.file_list),
                     test.name,
                 )
+
+    def test_get_wheels_from_manifest(self):
+        @dataclasses.dataclass
+        class TestParams:
+            name: str
+            toml_content: dict
+            expected: list
+
+        for test in [
+            TestParams(
+                name='no wheels',
+                toml_content={'type': 'add-on'},
+                expected=None,
+            ),
+            TestParams(
+                name='top-level wheels',
+                toml_content={
+                    'type': 'add-on',
+                    'wheels': ['./wheels/1.whl', './wheels/2.whl'],
+                },
+                expected=['./wheels/1.whl', './wheels/2.whl'],
+            ),
+            TestParams(
+                name='build.generated wheels',
+                toml_content={
+                    'type': 'add-on',
+                    'wheels': ['./wheels/1.whl', './wheels/2.whl'],
+                    'build': {'generated': {'wheels': ['./wheels/1.whl']}},
+                },
+                expected=['./wheels/1.whl'],
+            ),
+        ]:
+            with self.subTest(**dataclasses.asdict(test)):
+                self.assertEqual(
+                    test.expected,
+                    get_wheels_from_manifest(test.toml_content),
+                    test.name,
+                )
+
+    @patch(
+        'files.utils.get_wheel_sha256_from_pypi',
+        lambda _, __: ('blahblah', None),
+    )
+    def test_validate_wheels(self):
+        buff = io.BytesIO()
+        with tempfile.TemporaryDirectory() as output_dir:
+            test_file_path = os.path.join(output_dir, 'test_file.zip')
+            with zipfile.ZipFile(buff, mode='w') as file:
+                file.writestr('blender_manifest.toml', b'wheels = ["wheels/1.whl"]')
+                file.writestr('wheels/1.whl', b'')
+
+            with open(test_file_path, 'wb') as f:
+                f.write(buff.getvalue())
+
+            self.assertEqual(
+                validate_wheels(test_file_path, ['wheels/1.whl']).get('wheels/1.whl'),
+                'digest in archive=e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855'
+                ', digest on pypi=blahblah',
+            )