import unittest

from scripts.device import choose_runtime


GIB = 1024 ** 3


class DevicePolicyTests(unittest.TestCase):
    def test_eight_gigabyte_cuda_uses_group_offload(self):
        policy = choose_runtime('cuda', 8 * GIB)
        self.assertEqual(policy['offload'], 'group')
        self.assertFalse(policy['use_stream'])

    def test_large_cuda_keeps_a_whole_component_on_gpu(self):
        policy = choose_runtime('cuda', 32 * GIB)
        self.assertEqual(policy['offload'], 'model')

    def test_apple_silicon_stays_on_the_eager_resident_path(self):
        policy = choose_runtime('mps', 24 * GIB)
        self.assertEqual(policy['offload'], 'resident')
        self.assertFalse(policy['use_stream'])
        self.assertEqual(policy['dtype'], 'bfloat16')

    def test_cpu_is_resident(self):
        self.assertEqual(choose_runtime('cpu', 64 * GIB)['offload'], 'resident')

    def test_unknown_device_is_rejected(self):
        with self.assertRaises(ValueError):
            choose_runtime('xpu', 8 * GIB)


if __name__ == '__main__':
    unittest.main()