mirror of
https://github.com/tiennm99/litellm.git
synced 2026-07-12 19:04:59 +00:00
user
d853d3dcd4
The pre-existing s3:// / gcs:// custom-logger tests called ``get_instance_fn`` without ``config_file_path``, which means the new runtime gate (refuse remote URLs unless invoked from a config-file load) now raises ``ValueError`` before reaching the mocked download paths. Each test was exercising the documented startup config-file load scenario; pass ``config_file_path="/any/path"`` to make that intent explicit and route past the gate. Affected: test_s3_download_success, test_gcs_download_success, test_invalid_url_format, test_download_failure_handling, test_file_cleanup.
252 lines
9.8 KiB
Python
252 lines
9.8 KiB
Python
import pytest
|
|
import os
|
|
import tempfile
|
|
import importlib.util
|
|
from unittest.mock import patch, MagicMock
|
|
from litellm.proxy.types_utils.utils import (
|
|
get_instance_fn,
|
|
_load_instance_from_remote_storage,
|
|
)
|
|
|
|
|
|
class TestCustomLoggerS3GCS:
|
|
"""Test custom logger loading from S3/GCS using URL prefixes"""
|
|
|
|
@pytest.fixture
|
|
def sample_custom_logger_content(self):
|
|
"""Sample custom logger file content"""
|
|
return """
|
|
from litellm.integrations.custom_logger import CustomLogger
|
|
|
|
class TestCustomLogger(CustomLogger):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.initialized = True
|
|
|
|
def log_pre_api_call(self, model, messages, kwargs):
|
|
print(f"Pre-API call to {model}")
|
|
|
|
# Instance to be imported
|
|
test_logger_instance = TestCustomLogger()
|
|
"""
|
|
|
|
@pytest.fixture
|
|
def temp_config_dir(self):
|
|
"""Create a temporary directory for config files"""
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
yield temp_dir
|
|
|
|
def test_local_file_loading_still_works(
|
|
self, temp_config_dir, sample_custom_logger_content
|
|
):
|
|
"""Test that local file loading continues to work (no URL prefix)"""
|
|
# Create a local custom logger file
|
|
custom_logger_path = os.path.join(temp_config_dir, "test_custom_logger.py")
|
|
with open(custom_logger_path, "w") as f:
|
|
f.write(sample_custom_logger_content)
|
|
|
|
# Create a dummy config file
|
|
config_path = os.path.join(temp_config_dir, "config.yaml")
|
|
with open(config_path, "w") as f:
|
|
f.write("model_list: []")
|
|
|
|
# Test loading the custom logger (traditional way)
|
|
instance = get_instance_fn(
|
|
"test_custom_logger.test_logger_instance", config_path
|
|
)
|
|
|
|
assert instance is not None
|
|
assert hasattr(instance, "initialized")
|
|
assert instance.initialized is True
|
|
|
|
def test_s3_url_parsing(self):
|
|
"""Test S3 URL parsing"""
|
|
test_url = "s3://my-bucket/loggers/custom_callbacks.proxy_handler_instance"
|
|
|
|
# Mock the download function to avoid actual S3 calls
|
|
with patch(
|
|
"litellm.proxy.common_utils.load_config_utils.download_python_file_from_s3"
|
|
) as mock_download:
|
|
mock_download.return_value = (
|
|
False # Will cause failure, but we just want to test parsing
|
|
)
|
|
|
|
with pytest.raises(ImportError, match="Failed to download"):
|
|
_load_instance_from_remote_storage(test_url)
|
|
|
|
# Verify the download was called with correct parameters
|
|
mock_download.assert_called_once()
|
|
call_args = mock_download.call_args
|
|
assert call_args.kwargs["bucket_name"] == "my-bucket"
|
|
assert call_args.kwargs["object_key"] == "loggers/custom_callbacks.py"
|
|
|
|
def test_gcs_url_parsing(self):
|
|
"""Test GCS URL parsing"""
|
|
test_url = "gcs://my-bucket/custom_logger.my_instance"
|
|
|
|
# Mock the download function
|
|
with patch(
|
|
"litellm.proxy.types_utils.utils._download_gcs_file_wrapper"
|
|
) as mock_download:
|
|
mock_download.return_value = False # Will cause failure
|
|
|
|
with pytest.raises(ImportError, match="Failed to download"):
|
|
_load_instance_from_remote_storage(test_url)
|
|
|
|
# Verify the download was called with correct parameters
|
|
mock_download.assert_called_once()
|
|
call_args = mock_download.call_args
|
|
assert (
|
|
call_args[0][0] == "my-bucket"
|
|
) # bucket_name (positional for _download_gcs_file_wrapper)
|
|
assert call_args[0][1] == "custom_logger.py" # object_key
|
|
|
|
@patch("litellm.proxy.common_utils.load_config_utils.download_python_file_from_s3")
|
|
def test_s3_download_success(self, mock_s3_download, sample_custom_logger_content):
|
|
"""Test successful S3 download and loading"""
|
|
|
|
# Configure S3 download to succeed and create the file
|
|
def mock_download(bucket_name, object_key, local_file_path):
|
|
with open(local_file_path, "w") as f:
|
|
f.write(sample_custom_logger_content)
|
|
return True
|
|
|
|
mock_s3_download.side_effect = mock_download
|
|
|
|
# Test loading with S3 URL — pass config_file_path to indicate
|
|
# this is a startup config-file load (the documented operator
|
|
# flow that the runtime gate preserves).
|
|
test_url = "s3://test-bucket/test_custom_logger.test_logger_instance"
|
|
instance = get_instance_fn(test_url, config_file_path="/any/path")
|
|
|
|
assert instance is not None
|
|
assert hasattr(instance, "initialized")
|
|
assert instance.initialized is True
|
|
|
|
# Verify S3 download was called with correct parameters
|
|
mock_s3_download.assert_called_once()
|
|
call_args = mock_s3_download.call_args
|
|
assert call_args.kwargs["bucket_name"] == "test-bucket"
|
|
assert call_args.kwargs["object_key"] == "test_custom_logger.py"
|
|
|
|
@patch("litellm.proxy.types_utils.utils._download_gcs_file_wrapper")
|
|
def test_gcs_download_success(
|
|
self, mock_gcs_download, sample_custom_logger_content
|
|
):
|
|
"""Test successful GCS download and loading"""
|
|
|
|
# Configure GCS download to succeed and create the file
|
|
def mock_download(bucket_name, object_key, local_file_path):
|
|
with open(local_file_path, "w") as f:
|
|
f.write(sample_custom_logger_content)
|
|
return True
|
|
|
|
mock_gcs_download.side_effect = mock_download
|
|
|
|
# Test loading with GCS URL (startup config-file load path).
|
|
test_url = "gcs://test-bucket/test_custom_logger.test_logger_instance"
|
|
instance = get_instance_fn(test_url, config_file_path="/any/path")
|
|
|
|
assert instance is not None
|
|
assert hasattr(instance, "initialized")
|
|
assert instance.initialized is True
|
|
|
|
def test_nested_path_parsing(self):
|
|
"""Test parsing of nested paths in URLs"""
|
|
test_url = "s3://my-bucket/loggers/production/advanced_logger.handler_instance"
|
|
|
|
with patch(
|
|
"litellm.proxy.common_utils.load_config_utils.download_python_file_from_s3"
|
|
) as mock_download:
|
|
mock_download.return_value = False
|
|
|
|
with pytest.raises(ImportError):
|
|
_load_instance_from_remote_storage(test_url)
|
|
|
|
# Verify correct object key was generated
|
|
call_args = mock_download.call_args
|
|
assert (
|
|
call_args.kwargs["object_key"]
|
|
== "loggers/production/advanced_logger.py"
|
|
)
|
|
|
|
def test_invalid_url_schemes(self):
|
|
"""Test error handling for invalid URL schemes"""
|
|
# URLs that look like URLs but aren't s3:// or gcs:// will be treated as module names
|
|
# and fail with regular ImportError
|
|
with pytest.raises(ImportError):
|
|
get_instance_fn("http://bucket/module.instance")
|
|
|
|
with pytest.raises(ImportError):
|
|
get_instance_fn("ftp://bucket/module.instance")
|
|
|
|
def test_invalid_url_format(self):
|
|
"""Test error handling for invalid URL formats (config-file load path)."""
|
|
# Missing bucket
|
|
with pytest.raises(ImportError, match="Invalid URL format"):
|
|
get_instance_fn("s3://", config_file_path="/any/path")
|
|
|
|
# Missing path
|
|
with pytest.raises(ImportError, match="Invalid URL format"):
|
|
get_instance_fn("s3://bucket-only", config_file_path="/any/path")
|
|
|
|
# Missing instance name
|
|
with pytest.raises(ImportError, match="Invalid module specification"):
|
|
get_instance_fn("s3://bucket/module-only", config_file_path="/any/path")
|
|
|
|
# Including .py extension (common mistake)
|
|
with pytest.raises(
|
|
ImportError,
|
|
match="Don't include '\\.py' extension and you must specify the instance name",
|
|
):
|
|
get_instance_fn(
|
|
"s3://bucket/custom_guardrail.py", config_file_path="/any/path"
|
|
)
|
|
|
|
@patch("litellm.proxy.common_utils.load_config_utils.download_python_file_from_s3")
|
|
def test_download_failure_handling(self, mock_s3_download):
|
|
"""Test handling of download failures"""
|
|
mock_s3_download.return_value = False
|
|
|
|
test_url = "s3://test-bucket/failing_logger.instance"
|
|
|
|
with pytest.raises(ImportError, match="Failed to download"):
|
|
get_instance_fn(test_url, config_file_path="/any/path")
|
|
|
|
@patch("litellm.proxy.common_utils.load_config_utils.download_python_file_from_s3")
|
|
def test_file_cleanup(self, mock_s3_download, sample_custom_logger_content):
|
|
"""Test that temporary files are cleaned up"""
|
|
created_files = []
|
|
|
|
def mock_download(bucket_name, object_key, local_file_path):
|
|
created_files.append(local_file_path)
|
|
with open(local_file_path, "w") as f:
|
|
f.write(sample_custom_logger_content)
|
|
return True
|
|
|
|
mock_s3_download.side_effect = mock_download
|
|
|
|
test_url = "s3://test-bucket/test_custom_logger.test_logger_instance"
|
|
instance = get_instance_fn(test_url, config_file_path="/any/path")
|
|
|
|
assert instance is not None
|
|
|
|
# Verify file was created and then cleaned up
|
|
assert len(created_files) == 1
|
|
temp_file = created_files[0]
|
|
assert not os.path.exists(
|
|
temp_file
|
|
), f"Temporary file {temp_file} was not cleaned up"
|
|
|
|
def test_no_url_prefix_fallback(self, temp_config_dir):
|
|
"""Test fallback when no URL prefix is used and local file doesn't exist"""
|
|
config_path = os.path.join(temp_config_dir, "config.yaml")
|
|
with open(config_path, "w") as f:
|
|
f.write("model_list: []")
|
|
|
|
# Test that it tries local loading when no URL prefix is used
|
|
with pytest.raises(
|
|
ImportError, match="Could not import instance from nonexistent_logger"
|
|
):
|
|
get_instance_fn("nonexistent_logger.instance", config_path)
|