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