style: black-format db generate test files

Apply black (26.3.1, repo-pinned) to the new db generate test files so they
pass the lint job's `black --check .` after the rebase onto
litellm_internal_staging. Formatting only, no behavior change.
This commit is contained in:
Stephen Chin 2026-06-17 08:37:40 -07:00
parent bbfb5ad72e
commit 00e35d3019
4 changed files with 84 additions and 41 deletions

View file

@ -11,7 +11,6 @@ from unittest.mock import MagicMock, patch
from litellm_proxy_extras.utils import ProxyExtrasDBManager, _get_prisma_env
# ---------------------------------------------------------------------------
# Unit tests for _get_prisma_env() PATH injection
# ---------------------------------------------------------------------------
@ -23,15 +22,13 @@ def test_get_prisma_env_injects_scripts_dir_at_index_0(monkeypatch):
# Remove scripts dir from PATH so it's not already there
stripped_path = "/usr/local/bin:/usr/bin:/bin"
monkeypatch.setenv("PATH", stripped_path)
with (
patch("sysconfig.get_path", return_value=fake_scripts),
):
with (patch("sysconfig.get_path", return_value=fake_scripts),):
env = _get_prisma_env()
path_entries = env["PATH"].split(os.pathsep)
assert path_entries[0] == fake_scripts, (
f"Expected scripts dir at index 0, got: {env['PATH']}"
)
assert (
path_entries[0] == fake_scripts
), f"Expected scripts dir at index 0, got: {env['PATH']}"
# Original PATH entries preserved after the injected dir
assert "/usr/local/bin" in path_entries
assert "/usr/bin" in path_entries
@ -59,7 +56,9 @@ def test_get_prisma_env_handles_empty_path(monkeypatch):
env = _get_prisma_env()
assert env["PATH"] == fake_scripts
assert os.pathsep not in env["PATH"] or env["PATH"].rstrip(os.pathsep) == fake_scripts
assert (
os.pathsep not in env["PATH"] or env["PATH"].rstrip(os.pathsep) == fake_scripts
)
def test_get_prisma_env_falls_back_to_dirname_when_sysconfig_falsy(monkeypatch):
@ -129,14 +128,18 @@ def test_db_push_subprocess_receives_scripts_dir_on_path(monkeypatch, tmp_path):
with (
patch("sysconfig.get_path", return_value=fake_scripts),
patch.object(ProxyExtrasDBManager, "_get_prisma_dir", return_value=str(tmp_path)),
patch.object(
ProxyExtrasDBManager, "_get_prisma_dir", return_value=str(tmp_path)
),
patch("litellm_proxy_extras.utils._get_prisma_command", return_value="prisma"),
patch("subprocess.run", side_effect=fake_subprocess_run),
):
result = ProxyExtrasDBManager.setup_database(use_migrate=False)
assert result is True, "setup_database should return True on success"
assert len(captured_calls) == 1, f"Expected 1 subprocess call, got {len(captured_calls)}"
assert (
len(captured_calls) == 1
), f"Expected 1 subprocess call, got {len(captured_calls)}"
received_env = captured_calls[0]["env"]
assert received_env is not None, "subprocess.run must be called with env="

View file

@ -20,7 +20,9 @@ import pytest
@pytest.fixture
def non_activated_env():
venv_root = sysconfig.get_path("data") or os.path.dirname(sysconfig.get_path("scripts") or "")
venv_root = sysconfig.get_path("data") or os.path.dirname(
sysconfig.get_path("scripts") or ""
)
env = os.environ.copy()
env.pop("VIRTUAL_ENV", None)
env.pop("VIRTUAL_ENV_PROMPT", None)
@ -36,7 +38,7 @@ def fake_prisma_bin(tmp_path):
fake_bin = tmp_path / "fake_bin"
fake_bin.mkdir()
prisma = fake_bin / "prisma"
prisma.write_text("#!/bin/sh\necho \"fake-prisma called: $@\"\nexit 0\n")
prisma.write_text('#!/bin/sh\necho "fake-prisma called: $@"\nexit 0\n')
prisma.chmod(prisma.stat().st_mode | stat.S_IEXEC | stat.S_IXGRP | stat.S_IXOTH)
return fake_bin

View file

@ -39,8 +39,13 @@ def test_db_generate_success(cli_runner):
"litellm.proxy.client.cli.commands.db.PrismaManager._get_prisma_dir",
return_value="/fake/prisma/dir",
),
patch("litellm.proxy.client.cli.commands.db._get_prisma_command", return_value="prisma"),
patch("litellm.proxy.client.cli.commands.db._get_prisma_env", return_value=None),
patch(
"litellm.proxy.client.cli.commands.db._get_prisma_command",
return_value="prisma",
),
patch(
"litellm.proxy.client.cli.commands.db._get_prisma_env", return_value=None
),
patch("litellm.proxy.client.cli.commands.db.os.path.exists", return_value=True),
patch("litellm.proxy.client.cli.commands.db.subprocess.run", mock_run),
):
@ -63,8 +68,13 @@ def test_db_generate_schema_path_uses_get_prisma_dir(cli_runner):
"litellm.proxy.client.cli.commands.db.PrismaManager._get_prisma_dir",
return_value="/custom/litellm/proxy/dir",
),
patch("litellm.proxy.client.cli.commands.db._get_prisma_command", return_value="prisma"),
patch("litellm.proxy.client.cli.commands.db._get_prisma_env", return_value=None),
patch(
"litellm.proxy.client.cli.commands.db._get_prisma_command",
return_value="prisma",
),
patch(
"litellm.proxy.client.cli.commands.db._get_prisma_env", return_value=None
),
patch("litellm.proxy.client.cli.commands.db.os.path.exists", return_value=True),
patch("litellm.proxy.client.cli.commands.db.subprocess.run", mock_run),
):
@ -83,8 +93,13 @@ def test_db_generate_prisma_failure(cli_runner):
"litellm.proxy.client.cli.commands.db.PrismaManager._get_prisma_dir",
return_value="/fake/prisma/dir",
),
patch("litellm.proxy.client.cli.commands.db._get_prisma_command", return_value="prisma"),
patch("litellm.proxy.client.cli.commands.db._get_prisma_env", return_value=None),
patch(
"litellm.proxy.client.cli.commands.db._get_prisma_command",
return_value="prisma",
),
patch(
"litellm.proxy.client.cli.commands.db._get_prisma_env", return_value=None
),
patch("litellm.proxy.client.cli.commands.db.os.path.exists", return_value=True),
patch("litellm.proxy.client.cli.commands.db.subprocess.run", mock_run),
):
@ -100,9 +115,16 @@ def test_db_generate_schema_missing(cli_runner):
"litellm.proxy.client.cli.commands.db.PrismaManager._get_prisma_dir",
return_value="/fake/prisma/dir",
),
patch("litellm.proxy.client.cli.commands.db._get_prisma_command", return_value="prisma"),
patch("litellm.proxy.client.cli.commands.db._get_prisma_env", return_value=None),
patch("litellm.proxy.client.cli.commands.db.os.path.exists", return_value=False),
patch(
"litellm.proxy.client.cli.commands.db._get_prisma_command",
return_value="prisma",
),
patch(
"litellm.proxy.client.cli.commands.db._get_prisma_env", return_value=None
),
patch(
"litellm.proxy.client.cli.commands.db.os.path.exists", return_value=False
),
):
result = cli_runner.invoke(cli, ["db", "generate"])
@ -119,7 +141,9 @@ def test_db_generate_proxy_extras_not_installed(cli_runner):
},
):
# Patch the module-level flag directly since the import already happened
with patch("litellm.proxy.client.cli.commands.db._PROXY_EXTRAS_AVAILABLE", False):
with patch(
"litellm.proxy.client.cli.commands.db._PROXY_EXTRAS_AVAILABLE", False
):
result = cli_runner.invoke(cli, ["db", "generate"])
assert result.exit_code != 0
@ -179,7 +203,9 @@ def test_db_generate_env_includes_scripts_dir_on_path(cli_runner):
return_value=injected_env,
),
patch("litellm.proxy.client.cli.commands.db.os.path.exists", return_value=True),
patch("litellm.proxy.client.cli.commands.db.subprocess.run", side_effect=_capture),
patch(
"litellm.proxy.client.cli.commands.db.subprocess.run", side_effect=_capture
),
):
result = cli_runner.invoke(cli, ["db", "generate"])
@ -225,7 +251,9 @@ def test_db_generate_env_does_not_duplicate_scripts_dir(cli_runner):
return_value={"PATH": scripts_dir + os.pathsep + "/usr/bin"},
),
patch("litellm.proxy.client.cli.commands.db.os.path.exists", return_value=True),
patch("litellm.proxy.client.cli.commands.db.subprocess.run", side_effect=_capture),
patch(
"litellm.proxy.client.cli.commands.db.subprocess.run", side_effect=_capture
),
):
result = cli_runner.invoke(cli, ["db", "generate"])

View file

@ -7,7 +7,10 @@ from unittest.mock import patch
import pytest
from litellm.proxy.client.cli.commands.db import _get_generate_env, _get_venv_scripts_dir
from litellm.proxy.client.cli.commands.db import (
_get_generate_env,
_get_venv_scripts_dir,
)
_DB_MODULE = "litellm.proxy.client.cli.commands.db"
@ -100,15 +103,20 @@ def test_prisma_client_py_resolvable_via_path():
fake_exe = os.path.join(tmpdir, "prisma-client-py")
with open(fake_exe, "w") as f:
f.write("#!/bin/sh\necho fake\n")
os.chmod(fake_exe, os.stat(fake_exe).st_mode | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH)
os.chmod(
fake_exe,
os.stat(fake_exe).st_mode | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH,
)
with patch(f"{_DB_MODULE}._get_prisma_env", return_value={"PATH": f"{tmpdir}:/usr/bin"}):
with patch(
f"{_DB_MODULE}._get_prisma_env", return_value={"PATH": f"{tmpdir}:/usr/bin"}
):
env = _get_generate_env()
path_entries = env["PATH"].split(os.pathsep)
assert path_entries[0] == tmpdir, (
f"scripts dir must be first on PATH; got {path_entries!r}"
)
assert (
path_entries[0] == tmpdir
), f"scripts dir must be first on PATH; got {path_entries!r}"
found = shutil.which("prisma-client-py", path=env["PATH"])
assert found is not None, "prisma-client-py not findable via PATH lookup"
assert os.path.dirname(found) == tmpdir
@ -125,12 +133,14 @@ def test_prisma_client_py_resolvable_via_path():
def test_scripts_dir_promoted_to_index_zero_when_at_later_position():
fake_scripts = "/home/user/venv/bin"
promoted_path = os.pathsep.join([
fake_scripts,
"/usr/bin",
"/opt/other",
"/bin",
])
promoted_path = os.pathsep.join(
[
fake_scripts,
"/usr/bin",
"/opt/other",
"/bin",
]
)
promoted_env = {"PATH": promoted_path}
with patch(f"{_DB_MODULE}._get_prisma_env", return_value=promoted_env):
@ -143,9 +153,7 @@ def test_scripts_dir_promoted_to_index_zero_when_at_later_position():
f"_get_prisma_env() is responsible for the promotion; "
f"_get_generate_env() must faithfully delegate."
)
assert path_entries.count(fake_scripts) == 1, (
"scripts_dir must not be duplicated"
)
assert path_entries.count(fake_scripts) == 1, "scripts_dir must not be duplicated"
# ---------------------------------------------------------------------------
@ -167,7 +175,10 @@ def test_get_generate_env_delegation_chain_returns_scripts_dir_at_index_0():
fake_scripts = "/fake/venv/bin"
base_env = {"PATH": "/usr/bin:/bin"}
with (
patch(f"{_DB_MODULE}._get_prisma_env", return_value={**base_env, "PATH": f"{fake_scripts}:/usr/bin:/bin"}),
patch(
f"{_DB_MODULE}._get_prisma_env",
return_value={**base_env, "PATH": f"{fake_scripts}:/usr/bin:/bin"},
),
):
env = _get_generate_env()
@ -190,4 +201,3 @@ def test_get_generate_env_falls_back_to_os_environ_when_prisma_env_is_none():
# Must return a dict (a copy of os.environ) without raising
assert isinstance(env, dict)
assert len(env) > 0