mirror of
https://github.com/firestar5683/StarPilot.git
synced 2026-10-28 00:53:45 +08:00
firestar5683
30 lines
1.7 KiB
Python
30 lines
1.7 KiB
Python
import unittest, hashlib, io, pathlib
|
|
from unittest.mock import patch
|
|
from tinygrad.helpers import fetch_fw
|
|
|
|
class TestFirmwareFetch(unittest.TestCase):
|
|
def test_local_zstd_unknown_content_size_and_hash(self):
|
|
from zstandard import ZstdCompressor
|
|
firmware = b"verified AMD firmware fixture" * 20
|
|
digest = hashlib.sha256(firmware).hexdigest()
|
|
for write_content_size in (False, True):
|
|
compressed = ZstdCompressor(write_content_size=write_content_size).compress(firmware)
|
|
with self.subTest(write_content_size=write_content_size), patch.object(pathlib.Path, "is_file", return_value=True), \
|
|
patch.object(pathlib.Path, "open", return_value=io.BytesIO(compressed)), \
|
|
patch("tinygrad.helpers.fetch", side_effect=AssertionError("network fallback")):
|
|
self.assertEqual(fetch_fw("amdgpu", "fixture.bin", digest), firmware)
|
|
|
|
def test_local_zstd_wrong_hash_uses_pinned_fallback(self):
|
|
from zstandard import ZstdCompressor
|
|
compressed = ZstdCompressor(write_content_size=False).compress(b"wrong firmware")
|
|
with patch.object(pathlib.Path, "is_file", return_value=True), patch.object(pathlib.Path, "open", return_value=io.BytesIO(compressed)), \
|
|
patch("tinygrad.helpers.fetch", side_effect=RuntimeError("hash mismatch fallback")) as fallback:
|
|
with self.assertRaisesRegex(RuntimeError, "hash mismatch fallback"):
|
|
fetch_fw("amdgpu", "fixture.bin", hashlib.sha256(b"expected firmware").hexdigest())
|
|
fallback.assert_called_once_with(
|
|
"https://gitlab.com/kernel-firmware/linux-firmware/-/raw/0a6871b19abf5d6e024b5d208b101ae53e7fa0de/amdgpu/fixture.bin",
|
|
subdir="fw", sha256=hashlib.sha256(b"expected firmware").hexdigest())
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|