mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-09 03:18:44 +00:00
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:
parent
bbfb5ad72e
commit
00e35d3019
4 changed files with 84 additions and 41 deletions
|
|
@ -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="
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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"])
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue