File size: 2,085 Bytes
e9f9fd3 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 |
import unittest
from unittest.mock import patch, MagicMock
import os
import sys
# Add parent directory to path to import deoldify
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from deoldify.device_id import DeviceId
from deoldify._device import _Device
class TestDevice(unittest.TestCase):
def setUp(self):
self.device_manager = _Device()
@patch('torch.cuda.is_available')
def test_init_device_cpu(self, mock_cuda_available):
mock_cuda_available.return_value = False
# Mock xpu availability if it exists in the environment
with patch('torch.xpu.is_available', create=True) as mock_xpu_available:
mock_xpu_available.return_value = False
self.device_manager._init_device()
self.assertEqual(self.device_manager._backend, 'cpu')
@patch('torch.cuda.is_available')
def test_init_device_cuda(self, mock_cuda_available):
mock_cuda_available.return_value = True
# Ensure xpu is false so we fall through to cuda
with patch('torch.xpu.is_available', create=True) as mock_xpu_available:
mock_xpu_available.return_value = False
self.device_manager._init_device()
self.assertEqual(self.device_manager._backend, 'cuda')
def test_set_cpu(self):
self.device_manager.set(DeviceId.CPU)
self.assertEqual(self.device_manager.current(), DeviceId.CPU)
self.assertEqual(os.environ.get('CUDA_VISIBLE_DEVICES'), '')
@patch('torch.cuda.is_available')
def test_set_gpu_cuda(self, mock_cuda_available):
mock_cuda_available.return_value = True
with patch('torch.xpu.is_available', create=True) as mock_xpu_available:
mock_xpu_available.return_value = False
self.device_manager._init_device()
self.device_manager.set(DeviceId.GPU0)
self.assertEqual(self.device_manager.current(), DeviceId.GPU0)
self.assertEqual(os.environ.get('CUDA_VISIBLE_DEVICES'), '0')
if __name__ == '__main__':
unittest.main()
|