Files
searchi/tools/embeddinggemma2/test_provision_vulkan.py
T

145 lines
7.3 KiB
Python
Raw Normal View History

import hashlib
import io
from pathlib import Path
import tempfile
import subprocess
import unittest
import zipfile
from unittest.mock import patch
import provision_vulkan as provisioning
from provision_vulkan import download, sdk_verified, write_sdk_provenance, stage_loader
class DownloadTests(unittest.TestCase):
@staticmethod
def fake_sdk(sdk):
for name in ["Bin/glslc.exe", "Lib/vulkan-1.lib", "Include/vulkan/vulkan.h",
"Include/vulkan/vulkan_core.h"]:
path = sdk / name
path.parent.mkdir(parents=True, exist_ok=True)
path.write_bytes(name.encode())
def test_fresh_install_preserves_previous_cache_and_uses_copy_only(self):
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
sdk = root / "sdk"
sdk.mkdir()
(sdk / "old-marker").write_bytes(b"previous tree")
def install(command, **options):
fresh = Path(command[command.index("--root") + 1])
self.assertFalse(fresh.exists())
self.assertIn("copy_only=1", command)
self.assertEqual(options["env"]["QT_OPENGL"], "software")
self.assertEqual(options["env"]["SEARCHI_FORCE_CPU"], "1")
self.fake_sdk(fresh)
with patch("provision_vulkan.subprocess.run", side_effect=install):
provisioning.install_sdk(root, sdk, root / "installer.exe")
self.assertTrue(sdk_verified(sdk))
backups = list(root.glob("*-backup-*"))
self.assertEqual(len(backups), 1)
self.assertEqual((backups[0] / "old-marker").read_bytes(), b"previous tree")
def test_failed_installer_and_failed_swap_preserve_previous_provenance(self):
for failure in ["installer", "swap"]:
with self.subTest(failure=failure), tempfile.TemporaryDirectory() as directory:
root = Path(directory)
sdk = root / "sdk"
self.fake_sdk(sdk)
write_sdk_provenance(sdk)
old_manifest = (sdk / "searchi-provenance.json").read_bytes()
def install(command, **options):
if failure == "installer":
raise subprocess.CalledProcessError(1, command)
self.fake_sdk(Path(command[command.index("--root") + 1]))
real_replace = provisioning.os.replace
def replace(source, destination):
if failure == "swap" and "-staging-" in Path(source).name and Path(destination) == sdk:
raise PermissionError("injected swap failure")
real_replace(source, destination)
with patch("provision_vulkan.subprocess.run", side_effect=install), \
patch("provision_vulkan.os.replace", side_effect=replace):
with self.assertRaises((OSError, subprocess.CalledProcessError)):
provisioning.install_sdk(root, sdk, root / "installer.exe")
self.assertEqual((sdk / "searchi-provenance.json").read_bytes(), old_manifest)
self.assertTrue(sdk_verified(sdk))
def test_install_rejects_target_outside_owned_directory(self):
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
owned = root / "owned"
owned.mkdir()
with patch("provision_vulkan.subprocess.run", side_effect=AssertionError("installer ran")):
with self.assertRaisesRegex(ValueError, "escapes owned"):
provisioning.install_sdk(owned, root / "outside", root / "installer.exe")
def test_sdk_requires_provenance_and_detects_header_corruption(self):
with tempfile.TemporaryDirectory() as directory:
sdk = Path(directory)
for name in ["Bin/glslc.exe", "Lib/vulkan-1.lib", "Include/vulkan/vulkan.h",
"Include/vulkan/vulkan_core.h", "Include/vk_video/video.h"]:
path = sdk / name
path.parent.mkdir(parents=True, exist_ok=True)
path.write_bytes(name.encode())
self.assertFalse(sdk_verified(sdk))
write_sdk_provenance(sdk)
self.assertTrue(sdk_verified(sdk))
(sdk / "Include/vk_video/video.h").write_bytes(b"corrupt")
self.assertFalse(sdk_verified(sdk))
def test_verified_loader_is_not_replaced_and_bad_archive_preserves_it(self):
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
archive = root / "runtime.zip"
prefix = f"VulkanRT-X64-{provisioning.VERSION}-Components/"
with zipfile.ZipFile(archive, "w") as package:
package.writestr(prefix + "x64/vulkan-1.dll", b"verified loader")
package.writestr(prefix + "VulkanRT-License.txt", b"license")
package.writestr("../outside", b"ignored")
checksum = provisioning.digest(archive)
loader = root / "loader"
with patch.object(provisioning, "LOADER_SHA", checksum):
stage_loader(archive, loader)
with patch("provision_vulkan.os.replace", side_effect=AssertionError("rewrote loader")):
stage_loader(archive, loader)
archive.write_bytes(b"corrupt")
with self.assertRaisesRegex(ValueError, "archive checksum"):
stage_loader(archive, loader)
self.assertEqual((loader / "vulkan-1.dll").read_bytes(), b"verified loader")
self.assertFalse((root / "outside").exists())
def test_unsupported_architecture_fails_before_filesystem_or_network(self):
with patch("provision_vulkan.platform.machine", return_value="ARM64"), \
patch("provision_vulkan.Path.mkdir", side_effect=AssertionError("filesystem")), \
patch("provision_vulkan.urllib.request.urlopen", side_effect=AssertionError("network")):
with self.assertRaisesRegex(RuntimeError, "Windows x64 only"):
provisioning.provision()
def test_verified_cache_needs_no_network(self):
with tempfile.TemporaryDirectory() as directory:
path = Path(directory) / "sdk.exe"
path.write_bytes(b"verified SDK")
checksum = hashlib.sha256(path.read_bytes()).hexdigest()
with patch("provision_vulkan.urllib.request.urlopen", side_effect=AssertionError("unexpected network")):
download(path, "sdk.exe", checksum, 100)
def test_corrupt_or_oversized_download_preserves_previous_file(self):
for payload, limit, reason in [(b"wrong", 100, "checksum"), (b"too large", 2, "size ceiling")]:
with self.subTest(reason=reason), tempfile.TemporaryDirectory() as directory:
path = Path(directory) / "sdk.exe"
path.write_bytes(b"previous cached file")
with patch("provision_vulkan.urllib.request.urlopen", return_value=io.BytesIO(payload)):
with self.assertRaisesRegex(ValueError, reason):
download(path, "sdk.exe", hashlib.sha256(b"expected").hexdigest(), limit)
self.assertEqual(path.read_bytes(), b"previous cached file")
self.assertFalse(path.with_suffix(".exe.partial").exists())
if __name__ == "__main__":
unittest.main()