27 lines
787 B
Python
27 lines
787 B
Python
import sys
|
|
import unittest
|
|
from pathlib import Path
|
|
|
|
|
|
MODULE_DIR = Path(__file__).resolve().parents[1]
|
|
sys.path.insert(0, str(MODULE_DIR))
|
|
|
|
import inspect_coreml
|
|
|
|
|
|
class CoreMLDeviceCategoryTests(unittest.TestCase):
|
|
def test_classifies_compute_device_types(self):
|
|
NeuralEngineDevice = type("MLNeuralEngineComputeDevice", (), {})
|
|
GPUDevice = type("MLGPUComputeDevice", (), {})
|
|
CPUDevice = type("MLCPUComputeDevice", (), {})
|
|
self.assertEqual(
|
|
"neuralEngine",
|
|
inspect_coreml._device_category(NeuralEngineDevice()),
|
|
)
|
|
self.assertEqual("gpu", inspect_coreml._device_category(GPUDevice()))
|
|
self.assertEqual("cpu", inspect_coreml._device_category(CPUDevice()))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|