diff --git a/litellm-proxy-extras/tests/test_prisma_env_path_injection.py b/litellm-proxy-extras/tests/test_prisma_env_path_injection.py index 2be045f6156..1d17d270eda 100644 --- a/litellm-proxy-extras/tests/test_prisma_env_path_injection.py +++ b/litellm-proxy-extras/tests/test_prisma_env_path_injection.py @@ -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=" diff --git a/tests/proxy_cli_integration_tests/test_db_generate_integration.py b/tests/proxy_cli_integration_tests/test_db_generate_integration.py index aa86e582e91..9488011b449 100644 --- a/tests/proxy_cli_integration_tests/test_db_generate_integration.py +++ b/tests/proxy_cli_integration_tests/test_db_generate_integration.py @@ -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 diff --git a/tests/test_litellm/proxy/client/cli/test_db_commands.py b/tests/test_litellm/proxy/client/cli/test_db_commands.py index 16ba30da2e5..c93952e6095 100644 --- a/tests/test_litellm/proxy/client/cli/test_db_commands.py +++ b/tests/test_litellm/proxy/client/cli/test_db_commands.py @@ -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"]) diff --git a/tests/test_litellm/proxy/client/cli/test_db_generate_path_verifier.py b/tests/test_litellm/proxy/client/cli/test_db_generate_path_verifier.py index 42173788aaf..022fc6bab68 100644 --- a/tests/test_litellm/proxy/client/cli/test_db_generate_path_verifier.py +++ b/tests/test_litellm/proxy/client/cli/test_db_generate_path_verifier.py @@ -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 -