Files
searchi/tools/embeddinggemma2/test_provision_vulkan.py
T

146 lines
7.3 KiB
Python

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()