Spaces:
Running
Running
File size: 5,494 Bytes
9f69f35 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 |
"""Unit tests for token validation in llm_factory.py."""
from unittest.mock import MagicMock, patch
import pytest
from src.utils.exceptions import ConfigurationError
class TestGetPydanticAiModelTokenValidation:
"""Tests for get_pydantic_ai_model function with token validation."""
@patch("src.utils.llm_factory.settings")
@patch("src.utils.hf_error_handler.log_token_info")
@patch("src.utils.hf_error_handler.validate_hf_token")
@patch("pydantic_ai.providers.huggingface.HuggingFaceProvider")
@patch("pydantic_ai.models.huggingface.HuggingFaceModel")
def test_validates_oauth_token(
self,
mock_model_class,
mock_provider_class,
mock_validate,
mock_log,
mock_settings,
) -> None:
"""Should validate and log OAuth token when provided."""
mock_settings.hf_token = None
mock_settings.huggingface_api_key = None
mock_settings.huggingface_model = "test-model"
mock_validate.return_value = (True, None)
mock_model_class.return_value = MagicMock()
from src.utils.llm_factory import get_pydantic_ai_model
get_pydantic_ai_model(oauth_token="hf_test_token")
# Should log token info
mock_log.assert_called_once_with("hf_test_token", context="get_pydantic_ai_model")
# Should validate token
mock_validate.assert_called_once_with("hf_test_token")
@patch("src.utils.llm_factory.settings")
@patch("src.utils.hf_error_handler.log_token_info")
@patch("src.utils.hf_error_handler.validate_hf_token")
@patch("src.utils.llm_factory.logger")
@patch("pydantic_ai.providers.huggingface.HuggingFaceProvider")
@patch("pydantic_ai.models.huggingface.HuggingFaceModel")
def test_warns_on_invalid_token(
self,
mock_model_class,
mock_provider_class,
mock_logger,
mock_validate,
mock_log,
mock_settings,
) -> None:
"""Should warn when token validation fails."""
mock_settings.hf_token = None
mock_settings.huggingface_api_key = None
mock_settings.huggingface_model = "test-model"
mock_validate.return_value = (False, "Token too short")
mock_model_class.return_value = MagicMock()
from src.utils.llm_factory import get_pydantic_ai_model
get_pydantic_ai_model(oauth_token="short")
# Should warn about invalid token
warning_calls = [
call
for call in mock_logger.warning.call_args_list
if "Token validation failed" in str(call)
]
assert len(warning_calls) > 0
@patch("src.utils.llm_factory.settings")
@patch("src.utils.hf_error_handler.log_token_info")
@patch("src.utils.hf_error_handler.validate_hf_token")
@patch("pydantic_ai.providers.huggingface.HuggingFaceProvider")
@patch("pydantic_ai.models.huggingface.HuggingFaceModel")
def test_uses_env_token_when_oauth_not_provided(
self,
mock_model_class,
mock_provider_class,
mock_validate,
mock_log,
mock_settings,
) -> None:
"""Should use environment token when OAuth token not provided."""
mock_settings.hf_token = "hf_env_token"
mock_settings.huggingface_api_key = None
mock_settings.huggingface_model = "test-model"
mock_validate.return_value = (True, None)
mock_model_class.return_value = MagicMock()
from src.utils.llm_factory import get_pydantic_ai_model
get_pydantic_ai_model(oauth_token=None)
# Should log and validate env token
mock_log.assert_called_once_with("hf_env_token", context="get_pydantic_ai_model")
mock_validate.assert_called_once_with("hf_env_token")
@patch("src.utils.llm_factory.settings")
def test_raises_when_no_token_available(self, mock_settings) -> None:
"""Should raise ConfigurationError when no token is available."""
mock_settings.hf_token = None
mock_settings.huggingface_api_key = None
from src.utils.llm_factory import get_pydantic_ai_model
with pytest.raises(ConfigurationError, match="HuggingFace token required"):
get_pydantic_ai_model(oauth_token=None)
@patch("src.utils.llm_factory.settings")
@patch("src.utils.hf_error_handler.log_token_info")
@patch("src.utils.hf_error_handler.validate_hf_token")
@patch("pydantic_ai.providers.huggingface.HuggingFaceProvider")
@patch("pydantic_ai.models.huggingface.HuggingFaceModel")
def test_oauth_token_priority_over_env(
self,
mock_model_class,
mock_provider_class,
mock_validate,
mock_log,
mock_settings,
) -> None:
"""Should prefer OAuth token over environment token."""
mock_settings.hf_token = "hf_env_token"
mock_settings.huggingface_api_key = None
mock_settings.huggingface_model = "test-model"
mock_validate.return_value = (True, None)
mock_model_class.return_value = MagicMock()
from src.utils.llm_factory import get_pydantic_ai_model
get_pydantic_ai_model(oauth_token="hf_oauth_token")
# Should use OAuth token, not env token
mock_log.assert_called_once_with("hf_oauth_token", context="get_pydantic_ai_model")
mock_validate.assert_called_once_with("hf_oauth_token")
|