diff --git a/studio/backend/tests/test_utils.py b/studio/backend/tests/test_utils.py index 3c33b33cb3..9ff457bb0b 100644 --- a/studio/backend/tests/test_utils.py +++ b/studio/backend/tests/test_utils.py @@ -285,7 +285,7 @@ class TestLogGpuMemory: def test_does_not_raise(self): log_gpu_memory("test") - def test_logs_gpu_info_when_available(self, caplog): + def test_logs_gpu_info_when_available(self, capsys): fake_info = { "available": True, "backend": "cuda", @@ -295,35 +295,27 @@ class TestLogGpuMemory: "utilization_pct": 12.5, "free_gb": 14.0, } - import structlog - from loggers import get_logger - with ( - patch( - "utils.hardware.hardware.get_gpu_memory_info", return_value = fake_info - ), - caplog.at_level(logging.INFO, logger = "utils.hardware.hardware"), + with patch( + "utils.hardware.hardware.get_gpu_memory_info", return_value = fake_info ): log_gpu_memory("unit-test") - assert "unit-test" in caplog.text - assert "CUDA" in caplog.text - assert "FakeGPU" in caplog.text + captured = capsys.readouterr().out + assert "unit-test" in captured + assert "CUDA" in captured + assert "FakeGPU" in captured - def test_logs_cpu_fallback_when_no_gpu(self, caplog): + def test_logs_cpu_fallback_when_no_gpu(self, capsys): fake_info = {"available": False, "backend": "cpu"} - import structlog - from loggers import get_logger - with ( - patch( - "utils.hardware.hardware.get_gpu_memory_info", return_value = fake_info - ), - caplog.at_level(logging.INFO, logger = "utils.hardware.hardware"), + with patch( + "utils.hardware.hardware.get_gpu_memory_info", return_value = fake_info ): log_gpu_memory("cpu-test") - assert "No GPU available" in caplog.text + captured = capsys.readouterr().out + assert "No GPU available" in captured # ========== format_error_message() ==========