mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-09 22:31:41 +00:00
6a0d03914c
* test: drop the cwd-relative sys.path.insert calls from the test suite
TQ003 stands at 1,077 across 1,058 files, and 1,015 of them are the same shape:
sys.path.insert(0, os.path.abspath("../..")) and its deeper siblings. The
argument resolves against the working directory rather than the file, so from
the repo root, where every job runs pytest, it inserts the directory two levels
above the checkout. It has never pointed at litellm. The package is installed
into the environment anyway, which is what actually makes the import work, and
what the rule's message has said all along.
Removing them leaves 1,634 imports of sys and os with no remaining reference,
and those go too, except where another test module imports the name back out of
the file. The rest of TQ003 is 62 call sites that resolve against __file__ or a
variable, which are a different question and are left alone.
Collection is identical either way: 45,871 tests and the same 51 pre-existing
collection errors before and after, and ruff reports no new undefined name.
* test: drop the duplicate imports the sys.path sweep exposed to F811
* test(pre-call-utils): restore the os import the new bedrock tests need
415 lines
13 KiB
Python
415 lines
13 KiB
Python
import json
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
|
|
import litellm
|
|
from litellm.integrations.bitbucket import BitBucketPromptManager
|
|
from litellm.integrations.bitbucket.bitbucket_client import _sanitize_file_path
|
|
|
|
|
|
@patch("litellm.integrations.bitbucket.bitbucket_prompt_manager.BitBucketClient")
|
|
def test_bitbucket_prompt_integration_with_litellm(mock_client_class):
|
|
"""Test BitBucket prompt integration with LiteLLM completion."""
|
|
# Mock the BitBucket client
|
|
mock_client = MagicMock()
|
|
mock_client.get_file_content.return_value = """---
|
|
model: gpt-4
|
|
temperature: 0.7
|
|
max_tokens: 150
|
|
---
|
|
System: You are a helpful assistant.
|
|
|
|
User: {{user_message}}"""
|
|
mock_client_class.return_value = mock_client
|
|
|
|
# Configure BitBucket
|
|
bitbucket_config = {
|
|
"workspace": "test-workspace",
|
|
"repository": "test-repo",
|
|
"access_token": "test-token",
|
|
}
|
|
|
|
# Set global BitBucket configuration
|
|
litellm.set_global_bitbucket_config(bitbucket_config)
|
|
|
|
# Test that the configuration was set
|
|
assert litellm.global_bitbucket_config == bitbucket_config
|
|
|
|
|
|
@patch("litellm.integrations.bitbucket.bitbucket_prompt_manager.BitBucketClient")
|
|
def test_bitbucket_prompt_manager_initialization(mock_client_class):
|
|
"""Test BitBucketPromptManager initialization."""
|
|
# Mock the BitBucket client
|
|
mock_client = MagicMock()
|
|
mock_client.get_file_content.return_value = """---
|
|
model: gpt-4
|
|
temperature: 0.7
|
|
---
|
|
Hello {{name}}!"""
|
|
mock_client_class.return_value = mock_client
|
|
|
|
config = {
|
|
"workspace": "test-workspace",
|
|
"repository": "test-repo",
|
|
"access_token": "test-token",
|
|
}
|
|
|
|
manager = BitBucketPromptManager(config, prompt_id="test_prompt")
|
|
|
|
# Should have loaded the prompt
|
|
assert "test_prompt" in manager.prompt_manager.prompts
|
|
template = manager.prompt_manager.prompts["test_prompt"]
|
|
assert template.model == "gpt-4"
|
|
assert template.temperature == 0.7
|
|
|
|
# Test rendering
|
|
rendered = manager.prompt_manager.render_template("test_prompt", {"name": "World"})
|
|
assert rendered == "Hello World!"
|
|
|
|
|
|
@patch("litellm.integrations.bitbucket.bitbucket_prompt_manager.BitBucketClient")
|
|
def test_bitbucket_prompt_manager_error_handling(mock_client_class):
|
|
"""Test BitBucketPromptManager error handling."""
|
|
# Mock the BitBucket client to raise an error
|
|
mock_client = MagicMock()
|
|
mock_client.get_file_content.side_effect = Exception("BitBucket API error")
|
|
mock_client_class.return_value = mock_client
|
|
|
|
config = {
|
|
"workspace": "test-workspace",
|
|
"repository": "test-repo",
|
|
"access_token": "test-token",
|
|
}
|
|
|
|
manager = BitBucketPromptManager(config, prompt_id="test_prompt")
|
|
|
|
with pytest.raises(
|
|
Exception, match="Failed to load prompt 'test_prompt' from BitBucket"
|
|
):
|
|
_ = manager.prompt_manager
|
|
|
|
|
|
def test_bitbucket_prompt_manager_config_validation():
|
|
"""Test BitBucketPromptManager configuration validation."""
|
|
# Test missing required fields - validation happens when prompt_manager is accessed
|
|
manager = BitBucketPromptManager({})
|
|
|
|
with pytest.raises(
|
|
ValueError, match="workspace, repository, and access_token are required"
|
|
):
|
|
_ = manager.prompt_manager
|
|
|
|
manager = BitBucketPromptManager({"workspace": "test"})
|
|
|
|
with pytest.raises(
|
|
ValueError, match="workspace, repository, and access_token are required"
|
|
):
|
|
_ = manager.prompt_manager
|
|
|
|
manager = BitBucketPromptManager({"repository": "test"})
|
|
|
|
with pytest.raises(
|
|
ValueError, match="workspace, repository, and access_token are required"
|
|
):
|
|
_ = manager.prompt_manager
|
|
|
|
manager = BitBucketPromptManager({"access_token": "test"})
|
|
|
|
with pytest.raises(
|
|
ValueError, match="workspace, repository, and access_token are required"
|
|
):
|
|
_ = manager.prompt_manager
|
|
|
|
|
|
@patch("litellm.integrations.bitbucket.bitbucket_prompt_manager.BitBucketClient")
|
|
def test_bitbucket_prompt_manager_complex_prompt(mock_client_class):
|
|
"""Test BitBucketPromptManager with complex prompt structure."""
|
|
# Mock the BitBucket client
|
|
mock_client = MagicMock()
|
|
mock_client.get_file_content.return_value = """---
|
|
model: gpt-4
|
|
temperature: 0.3
|
|
max_tokens: 500
|
|
input:
|
|
schema:
|
|
user_question: string
|
|
context?: string
|
|
language: string
|
|
---
|
|
System: You are a helpful {{language}} programming assistant.
|
|
|
|
{% if context %}Context: {{context}}
|
|
|
|
{% endif %}User: {{user_question}}
|
|
|
|
Please provide a detailed response in {{language}}."""
|
|
mock_client_class.return_value = mock_client
|
|
|
|
config = {
|
|
"workspace": "test-workspace",
|
|
"repository": "test-repo",
|
|
"access_token": "test-token",
|
|
}
|
|
|
|
manager = BitBucketPromptManager(config, prompt_id="complex_prompt")
|
|
|
|
# Should have loaded the prompt
|
|
assert "complex_prompt" in manager.prompt_manager.prompts
|
|
template = manager.prompt_manager.prompts["complex_prompt"]
|
|
assert template.model == "gpt-4"
|
|
assert template.temperature == 0.3
|
|
assert template.max_tokens == 500
|
|
assert template.input_schema == {
|
|
"user_question": "string",
|
|
"context?": "string",
|
|
"language": "string",
|
|
}
|
|
|
|
# Test rendering with all variables
|
|
rendered = manager.prompt_manager.render_template(
|
|
"complex_prompt",
|
|
{
|
|
"user_question": "How do I create a class?",
|
|
"context": "Python programming",
|
|
"language": "Python",
|
|
},
|
|
)
|
|
|
|
assert "You are a helpful Python programming assistant." in rendered
|
|
assert "Context: Python programming" in rendered
|
|
assert "How do I create a class?" in rendered
|
|
assert "Please provide a detailed response in Python." in rendered
|
|
|
|
# Test rendering without optional context
|
|
rendered_no_context = manager.prompt_manager.render_template(
|
|
"complex_prompt", {"user_question": "What is inheritance?", "language": "Java"}
|
|
)
|
|
|
|
assert "You are a helpful Java programming assistant." in rendered_no_context
|
|
assert "Context:" not in rendered_no_context
|
|
assert "What is inheritance?" in rendered_no_context
|
|
|
|
|
|
@patch("litellm.integrations.bitbucket.bitbucket_prompt_manager.BitBucketClient")
|
|
def test_bitbucket_prompt_manager_message_parsing(mock_client_class):
|
|
"""Test BitBucketPromptManager message parsing for different prompt formats."""
|
|
# Mock the BitBucket client
|
|
mock_client = MagicMock()
|
|
mock_client.get_file_content.return_value = """---
|
|
model: gpt-4
|
|
---
|
|
System: You are a helpful assistant.
|
|
|
|
User: {{user_message}}
|
|
|
|
Assistant: I'll help you with that."""
|
|
mock_client_class.return_value = mock_client
|
|
|
|
config = {
|
|
"workspace": "test-workspace",
|
|
"repository": "test-repo",
|
|
"access_token": "test-token",
|
|
}
|
|
|
|
manager = BitBucketPromptManager(config, prompt_id="conversation_prompt")
|
|
|
|
# Test message parsing
|
|
messages = manager._parse_prompt_to_messages(
|
|
"System: You are a helpful assistant.\n\nUser: Hello!\n\nAssistant: Hi there!"
|
|
)
|
|
|
|
assert len(messages) == 3
|
|
assert messages[0]["role"] == "system"
|
|
assert messages[0]["content"] == "You are a helpful assistant."
|
|
assert messages[1]["role"] == "user"
|
|
assert messages[1]["content"] == "Hello!"
|
|
assert messages[2]["role"] == "assistant"
|
|
assert messages[2]["content"] == "Hi there!"
|
|
|
|
|
|
@patch("litellm.integrations.bitbucket.bitbucket_prompt_manager.BitBucketClient")
|
|
def test_bitbucket_prompt_manager_pre_call_hook_integration(mock_client_class):
|
|
"""Test BitBucketPromptManager pre_call_hook integration."""
|
|
# Mock the BitBucket client
|
|
mock_client = MagicMock()
|
|
mock_client.get_file_content.return_value = """---
|
|
model: gpt-4
|
|
temperature: 0.8
|
|
max_tokens: 200
|
|
---
|
|
System: You are a helpful assistant.
|
|
|
|
User: {{user_message}}"""
|
|
mock_client_class.return_value = mock_client
|
|
|
|
config = {
|
|
"workspace": "test-workspace",
|
|
"repository": "test-repo",
|
|
"access_token": "test-token",
|
|
}
|
|
|
|
manager = BitBucketPromptManager(config, prompt_id="test_prompt")
|
|
|
|
# Test pre_call_hook
|
|
original_messages = [{"role": "user", "content": "This will be ignored"}]
|
|
litellm_params = {"api_key": "test-key"}
|
|
|
|
result_messages, result_params = manager.pre_call_hook(
|
|
user_id="test_user",
|
|
messages=original_messages,
|
|
litellm_params=litellm_params,
|
|
prompt_id="test_prompt",
|
|
prompt_variables={"user_message": "What is AI?"},
|
|
)
|
|
|
|
# Should have parsed the prompt into messages
|
|
assert len(result_messages) == 2
|
|
assert result_messages[0]["role"] == "system"
|
|
assert result_messages[0]["content"] == "You are a helpful assistant."
|
|
assert result_messages[1]["role"] == "user"
|
|
assert result_messages[1]["content"] == "What is AI?"
|
|
|
|
# Should have updated litellm_params
|
|
assert result_params["model"] == "gpt-4"
|
|
assert result_params["temperature"] == 0.8
|
|
assert result_params["max_tokens"] == 200
|
|
assert result_params["api_key"] == "test-key" # Original params preserved
|
|
|
|
|
|
@patch("litellm.integrations.bitbucket.bitbucket_prompt_manager.BitBucketClient")
|
|
def test_bitbucket_prompt_manager_post_call_hook(mock_client_class):
|
|
"""Test BitBucketPromptManager post_call_hook."""
|
|
# Mock the BitBucket client
|
|
mock_client = MagicMock()
|
|
mock_client.get_file_content.return_value = "Simple prompt: {{message}}"
|
|
mock_client_class.return_value = mock_client
|
|
|
|
config = {
|
|
"workspace": "test-workspace",
|
|
"repository": "test-repo",
|
|
"access_token": "test-token",
|
|
}
|
|
|
|
manager = BitBucketPromptManager(config, prompt_id="test_prompt")
|
|
|
|
# Mock response
|
|
mock_response = MagicMock()
|
|
mock_response.choices = [MagicMock()]
|
|
mock_response.choices[0].message = MagicMock()
|
|
mock_response.choices[0].message.content = "Test response"
|
|
|
|
# Test post_call_hook
|
|
result = manager.post_call_hook(
|
|
user_id="test_user",
|
|
response=mock_response,
|
|
input_messages=[{"role": "user", "content": "test"}],
|
|
litellm_params={},
|
|
prompt_id="test_prompt",
|
|
)
|
|
|
|
# Should return the response unchanged
|
|
assert result == mock_response
|
|
|
|
|
|
def test_bitbucket_prompt_manager_integration_name():
|
|
"""Test BitBucketPromptManager integration name."""
|
|
config = {
|
|
"workspace": "test-workspace",
|
|
"repository": "test-repo",
|
|
"access_token": "test-token",
|
|
}
|
|
|
|
manager = BitBucketPromptManager(config)
|
|
assert manager.integration_name == "bitbucket"
|
|
|
|
|
|
@patch("litellm.integrations.bitbucket.bitbucket_prompt_manager.BitBucketClient")
|
|
def test_bitbucket_prompt_manager_get_template(mock_client_class):
|
|
"""Test BitBucketPromptManager get_template method."""
|
|
# Mock the BitBucket client
|
|
mock_client = MagicMock()
|
|
mock_client.get_file_content.return_value = "Test content"
|
|
mock_client_class.return_value = mock_client
|
|
|
|
config = {
|
|
"workspace": "test-workspace",
|
|
"repository": "test-repo",
|
|
"access_token": "test-token",
|
|
}
|
|
|
|
manager = BitBucketPromptManager(config, prompt_id="test_prompt")
|
|
|
|
# Test getting existing template
|
|
template = manager.prompt_manager.get_template("test_prompt")
|
|
assert template is not None
|
|
assert template.template_id == "test_prompt"
|
|
|
|
# Test getting non-existing template
|
|
template = manager.prompt_manager.get_template("nonexistent")
|
|
assert template is None
|
|
|
|
|
|
@patch("litellm.integrations.bitbucket.bitbucket_prompt_manager.BitBucketClient")
|
|
def test_bitbucket_prompt_manager_list_templates(mock_client_class):
|
|
"""Test BitBucketPromptManager list_templates method."""
|
|
# Mock the BitBucket client
|
|
mock_client = MagicMock()
|
|
mock_client.get_file_content.return_value = "Test content"
|
|
mock_client_class.return_value = mock_client
|
|
|
|
config = {
|
|
"workspace": "test-workspace",
|
|
"repository": "test-repo",
|
|
"access_token": "test-token",
|
|
}
|
|
|
|
manager = BitBucketPromptManager(config, prompt_id="test_prompt")
|
|
|
|
# Test listing templates
|
|
templates = manager.prompt_manager.list_templates()
|
|
assert isinstance(templates, list)
|
|
assert "test_prompt" in templates
|
|
|
|
|
|
# --- Security: path traversal / SSRF ---
|
|
|
|
|
|
def test_sanitize_file_path_rejects_traversal():
|
|
with pytest.raises(ValueError, match="path traversal"):
|
|
_sanitize_file_path("../../etc/passwd")
|
|
|
|
|
|
def test_sanitize_file_path_rejects_fragment():
|
|
with pytest.raises(ValueError, match="URL special characters"):
|
|
_sanitize_file_path("secret#.prompt")
|
|
|
|
|
|
def test_sanitize_file_path_rejects_query():
|
|
with pytest.raises(ValueError, match="URL special characters"):
|
|
_sanitize_file_path("secret?.prompt")
|
|
|
|
|
|
def test_sanitize_file_path_encodes_special_chars():
|
|
result = _sanitize_file_path("prompts/my prompt.prompt")
|
|
assert result == "prompts/my%20prompt.prompt"
|
|
|
|
|
|
def test_sanitize_file_path_allows_normal_paths():
|
|
assert _sanitize_file_path("prompts/my-prompt") == "prompts/my-prompt"
|
|
assert _sanitize_file_path("simple") == "simple"
|
|
|
|
|
|
def test_bitbucket_client_rejects_traversal_in_get_file_content():
|
|
from litellm.integrations.bitbucket.bitbucket_client import BitBucketClient
|
|
|
|
client = BitBucketClient(
|
|
{
|
|
"workspace": "ws",
|
|
"repository": "repo",
|
|
"access_token": "tok",
|
|
}
|
|
)
|
|
with pytest.raises(ValueError, match="path traversal"):
|
|
client.get_file_content("../../admin/credentials")
|