mirror of
https://github.com/BerriAI/litellm.git
synced 2025-04-26 03:04:13 +00:00
55 lines
1.7 KiB
Python
55 lines
1.7 KiB
Python
"""
|
|
Test the ASI common utilities module
|
|
"""
|
|
|
|
import unittest
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
from litellm.llms.asi.common_utils import (
|
|
is_asi_model,
|
|
get_asi_model_name,
|
|
validate_environment,
|
|
)
|
|
|
|
|
|
class TestASICommonUtils(unittest.TestCase):
|
|
"""Test the ASI common utilities"""
|
|
|
|
def test_is_asi_model(self):
|
|
"""Test the is_asi_model function"""
|
|
# Test with ASI model
|
|
self.assertTrue(is_asi_model("asi1-mini"))
|
|
self.assertTrue(is_asi_model("asi/asi1-mini"))
|
|
|
|
# Test with non-ASI model
|
|
self.assertFalse(is_asi_model("gpt-4"))
|
|
self.assertFalse(is_asi_model("claude-3"))
|
|
|
|
def test_get_asi_model_name(self):
|
|
"""Test the get_asi_model_name function"""
|
|
# Test with provider prefix
|
|
self.assertEqual(get_asi_model_name("asi/asi1-mini"), "asi1-mini")
|
|
|
|
# Test without provider prefix
|
|
self.assertEqual(get_asi_model_name("asi1-mini"), "asi1-mini")
|
|
|
|
@patch("litellm.utils.get_secret")
|
|
def test_validate_environment(self, mock_get_secret):
|
|
"""Test the validate_environment function"""
|
|
# Test with provided API key
|
|
api_key = "test-api-key"
|
|
validate_environment(api_key) # Should not raise an error
|
|
|
|
# Test with environment variable
|
|
mock_get_secret.return_value = "env-api-key"
|
|
validate_environment() # Should not raise an error
|
|
mock_get_secret.assert_called_with("ASI_API_KEY")
|
|
|
|
# Test with missing API key
|
|
mock_get_secret.return_value = None
|
|
with self.assertRaises(ValueError):
|
|
validate_environment()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|