-
-
Notifications
You must be signed in to change notification settings - Fork 5.9k
Expand file tree
/
Copy pathtest_install_python_stack.py
More file actions
56 lines (44 loc) · 2.1 KB
/
test_install_python_stack.py
File metadata and controls
56 lines (44 loc) · 2.1 KB
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
53
54
55
56
"""Tests for install_python_stack._build_uv_cmd torch-backend handling."""
from __future__ import annotations
import importlib
import os
import sys
from pathlib import Path
from unittest import mock
import pytest
# Add the studio directory so we can import install_python_stack
STUDIO_DIR = Path(__file__).resolve().parents[2] / "studio"
sys.path.insert(0, str(STUDIO_DIR))
# _build_uv_cmd lives at module level; import after path setup.
# We need to mock parts of the module that do work at import time.
import install_python_stack as ips
class TestBuildUvCmdTorchBackend:
"""Verify _build_uv_cmd only adds --torch-backend when UV_TORCH_BACKEND is set."""
def _call(self, args: tuple[str, ...] = ()) -> list[str]:
return ips._build_uv_cmd(args)
def test_default_no_torch_backend(self):
"""Without UV_TORCH_BACKEND env var, no --torch-backend flag."""
env = os.environ.copy()
env.pop("UV_TORCH_BACKEND", None)
with mock.patch.dict(os.environ, env, clear = True):
cmd = self._call(("somepackage",))
assert not any(
a.startswith("--torch-backend") for a in cmd
), f"--torch-backend should not appear by default, got: {cmd}"
def test_uv_torch_backend_auto(self):
"""UV_TORCH_BACKEND=auto adds --torch-backend=auto."""
with mock.patch.dict(os.environ, {"UV_TORCH_BACKEND": "auto"}):
cmd = self._call(("somepackage",))
assert "--torch-backend=auto" in cmd
def test_uv_torch_backend_cpu(self):
"""UV_TORCH_BACKEND=cpu adds --torch-backend=cpu."""
with mock.patch.dict(os.environ, {"UV_TORCH_BACKEND": "cpu"}):
cmd = self._call(("somepackage",))
assert "--torch-backend=cpu" in cmd
def test_uv_torch_backend_empty(self):
"""UV_TORCH_BACKEND="" (empty string) should NOT add --torch-backend."""
with mock.patch.dict(os.environ, {"UV_TORCH_BACKEND": ""}):
cmd = self._call(("somepackage",))
assert not any(
a.startswith("--torch-backend") for a in cmd
), f"Empty UV_TORCH_BACKEND should not add flag, got: {cmd}"