mirror of
https://github.com/firestar5683/StarPilot.git
synced 2026-08-01 14:16:45 +08:00
firestar5683
39 lines
1.4 KiB
Python
39 lines
1.4 KiB
Python
import hashlib
|
|
from pathlib import Path
|
|
import tempfile
|
|
import unittest
|
|
from unittest import mock
|
|
|
|
import numpy as np
|
|
import zstandard
|
|
import tinygrad.helpers
|
|
from tinygrad.helpers import _decompress_zstd, fetch_fw, polyN, is_numpy_ndarray
|
|
from tinygrad.tensor import Tensor
|
|
|
|
class TestPolyN(unittest.TestCase):
|
|
def test_tensor(self):
|
|
np.testing.assert_allclose(polyN(Tensor([1.0, 2.0, 3.0, 4.0]), [1.0, -2.0, 1.0]).numpy(), [0.0, 1.0, 4.0, 9.0])
|
|
|
|
class TestIsNumpyNdarray(unittest.TestCase):
|
|
def test_tensor_numpy(self):
|
|
self.assertTrue(is_numpy_ndarray(Tensor([1, 2, 3]).numpy()))
|
|
|
|
class TestZstd(unittest.TestCase):
|
|
def test_decompress(self):
|
|
payload = b"local firmware payload"
|
|
compressed = zstandard.ZstdCompressor().compress(payload)
|
|
self.assertEqual(_decompress_zstd(compressed), payload)
|
|
|
|
def test_fetch_fw_uses_local_zstd(self):
|
|
payload = b"local firmware payload"
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
firmware = Path(tmp) / "firmware.bin.zst"
|
|
firmware.write_bytes(zstandard.ZstdCompressor().compress(payload))
|
|
with mock.patch.object(tinygrad.helpers.pathlib, "Path", return_value=firmware), \
|
|
mock.patch.object(tinygrad.helpers, "fetch") as remote_fetch:
|
|
self.assertEqual(fetch_fw("amdgpu", "firmware.bin", hashlib.sha256(payload).hexdigest()), payload)
|
|
remote_fetch.assert_not_called()
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|