File size: 1,590 Bytes
46728df | 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 | import os
from unittest.mock import patch
from server.device import select_device
def test_env_override_cuda():
with patch.dict(os.environ, {"CHATTERBOX_DEVICE": "cuda"}):
assert select_device() == "cuda"
def test_env_override_mps():
with patch.dict(os.environ, {"CHATTERBOX_DEVICE": "MPS"}):
assert select_device() == "mps"
def test_env_override_cpu():
with patch.dict(os.environ, {"CHATTERBOX_DEVICE": "cpu"}):
assert select_device() == "cpu"
def test_invalid_env_falls_through_to_autodetect():
with patch.dict(os.environ, {"CHATTERBOX_DEVICE": "tpu"}, clear=False):
with patch("server.device._cuda_available", return_value=True):
assert select_device() == "cuda"
def test_autodetect_prefers_cuda_over_mps():
with patch.dict(os.environ, {}, clear=True):
with patch("server.device._cuda_available", return_value=True), \
patch("server.device._mps_available", return_value=True):
assert select_device() == "cuda"
def test_autodetect_uses_mps_when_no_cuda():
with patch.dict(os.environ, {}, clear=True):
with patch("server.device._cuda_available", return_value=False), \
patch("server.device._mps_available", return_value=True):
assert select_device() == "mps"
def test_autodetect_falls_back_to_cpu():
with patch.dict(os.environ, {}, clear=True):
with patch("server.device._cuda_available", return_value=False), \
patch("server.device._mps_available", return_value=False):
assert select_device() == "cpu"
|