Spaces:
Paused
Paused
File size: 14,200 Bytes
a5784e9 | 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 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 | """
Tests for api_utils/context_init.py - Request context initialization.
Test Strategy:
- Test initialize_request_context with various configurations
- Use REAL asyncio.Lock instances (not AsyncMock)
- Use real server_state module
- Mock only logging side effects
- Test context dictionary construction with actual values
Coverage Target: 95%+
Mock Budget: <40 (down from 95)
"""
import asyncio
from unittest.mock import MagicMock, patch
import pytest
from api_utils.context_init import initialize_request_context
from api_utils.server_state import state
class TestInitializeRequestContext:
"""Tests for initialize_request_context function."""
@pytest.mark.asyncio
async def test_streaming_request_context(
self, real_locks_mock_browser, make_chat_request
):
"""Test context initialization for streaming request."""
# Set up server state
state.current_ai_studio_model_id = "gemini-1.5-pro"
state.is_page_ready = True
state.page_instance = real_locks_mock_browser.page_instance
state.model_switching_lock = real_locks_mock_browser.model_switching_lock
state.params_cache_lock = real_locks_mock_browser.params_cache_lock
state.page_params_cache = {"temperature": 1.0}
state.parsed_model_list = [
{"id": "gemini-1.5-pro", "object": "model"},
{"id": "gemini-1.5-flash", "object": "model"},
]
request = make_chat_request(model="gemini-1.5-pro", stream=True)
with patch("api_utils.server_state.state.logger", MagicMock()) as mock_logger:
context = await initialize_request_context("req1", request)
# Verify logging - now uses debug, not info
assert mock_logger.debug.call_count >= 1
log_calls = [call[0][0] for call in mock_logger.debug.call_args_list]
assert any("[Request]" in msg for msg in log_calls)
# Verify context fields
assert context["is_streaming"] is True
assert context["requested_model"] == "gemini-1.5-pro"
assert context["current_ai_studio_model_id"] == "gemini-1.5-pro"
assert context["is_page_ready"] is True
assert context["page"] == real_locks_mock_browser.page_instance
assert context["model_actually_switched"] is False
assert context["needs_model_switching"] is False
assert context["model_id_to_use"] is None
# Verify real locks
assert isinstance(context["model_switching_lock"], asyncio.Lock)
assert isinstance(context["params_cache_lock"], asyncio.Lock)
@pytest.mark.asyncio
async def test_non_streaming_request_context(
self, real_locks_mock_browser, make_chat_request
):
"""Test context initialization for non-streaming request."""
state.current_ai_studio_model_id = "gemini-1.5-flash"
state.is_page_ready = True
state.page_instance = real_locks_mock_browser.page_instance
state.model_switching_lock = real_locks_mock_browser.model_switching_lock
state.params_cache_lock = real_locks_mock_browser.params_cache_lock
state.page_params_cache = {}
state.parsed_model_list = []
request = make_chat_request(model="gemini-1.5-flash", stream=False)
with patch("api_utils.server_state.state.logger", MagicMock()) as mock_logger:
context = await initialize_request_context("req2", request)
# Verify streaming flag
assert context["is_streaming"] is False
assert context["requested_model"] == "gemini-1.5-flash"
# Verify logging uses debug
log_calls = [call[0][0] for call in mock_logger.debug.call_args_list]
assert any("Stream=False" in msg for msg in log_calls)
@pytest.mark.asyncio
async def test_different_model_name(
self, real_locks_mock_browser, make_chat_request
):
"""Test context with different model name."""
state.current_ai_studio_model_id = "gemini-1.5-pro"
state.is_page_ready = True
state.page_instance = real_locks_mock_browser.page_instance
state.model_switching_lock = real_locks_mock_browser.model_switching_lock
state.params_cache_lock = real_locks_mock_browser.params_cache_lock
state.page_params_cache = {}
state.parsed_model_list = []
request = make_chat_request(model="gemini-2.0-flash-thinking-exp")
with patch("api_utils.server_state.state.logger", MagicMock()) as mock_logger:
context = await initialize_request_context("req3", request)
assert context["requested_model"] == "gemini-2.0-flash-thinking-exp"
# Verify logging includes model name in debug call
log_calls = [call[0][0] for call in mock_logger.debug.call_args_list]
assert any("gemini-2.0-flash-thinking-exp" in msg for msg in log_calls)
@pytest.mark.asyncio
@pytest.mark.parametrize(
"is_ready,expected",
[
(True, True),
(False, False),
],
)
async def test_page_ready_states(
self, real_locks_mock_browser, make_chat_request, is_ready, expected
):
"""Test context with different page ready states."""
state.current_ai_studio_model_id = "gemini-1.5-pro"
state.is_page_ready = is_ready
state.page_instance = real_locks_mock_browser.page_instance
state.model_switching_lock = real_locks_mock_browser.model_switching_lock
state.params_cache_lock = real_locks_mock_browser.params_cache_lock
state.page_params_cache = {}
state.parsed_model_list = []
request = make_chat_request()
with patch("api_utils.server_state.state.logger", MagicMock()):
context = await initialize_request_context("req4", request)
assert context["is_page_ready"] == expected
@pytest.mark.asyncio
async def test_none_current_model(self, real_locks_mock_browser, make_chat_request):
"""Test context when current model ID is None (initial state)."""
state.current_ai_studio_model_id = None
state.is_page_ready = True
state.page_instance = real_locks_mock_browser.page_instance
state.model_switching_lock = real_locks_mock_browser.model_switching_lock
state.params_cache_lock = real_locks_mock_browser.params_cache_lock
state.page_params_cache = {}
state.parsed_model_list = []
request = make_chat_request()
with patch("api_utils.server_state.state.logger", MagicMock()):
context = await initialize_request_context("req5", request)
assert context["current_ai_studio_model_id"] is None
@pytest.mark.asyncio
@pytest.mark.parametrize(
"cache_value",
[
{},
{"temperature": 0.7, "max_tokens": 1024},
{"last_known_model_id_for_params": "gemini-1.5-pro"},
],
)
async def test_various_params_caches(
self, real_locks_mock_browser, make_chat_request, cache_value
):
"""Test context with various parameter cache states."""
state.current_ai_studio_model_id = "gemini-1.5-pro"
state.is_page_ready = True
state.page_instance = real_locks_mock_browser.page_instance
state.model_switching_lock = real_locks_mock_browser.model_switching_lock
state.params_cache_lock = real_locks_mock_browser.params_cache_lock
state.page_params_cache = cache_value
state.parsed_model_list = []
request = make_chat_request()
with patch("api_utils.server_state.state.logger", MagicMock()):
context = await initialize_request_context("req6", request)
assert context["page_params_cache"] == cache_value
@pytest.mark.asyncio
@pytest.mark.parametrize(
"model_list",
[
[],
[{"id": "gemini-1.5-pro", "object": "model"}],
[
{"id": "gemini-1.5-pro", "object": "model"},
{"id": "gemini-1.5-flash", "object": "model"},
{"id": "gemini-2.0-flash", "object": "model"},
],
],
)
async def test_various_model_lists(
self, real_locks_mock_browser, make_chat_request, model_list
):
"""Test context with various model lists."""
state.current_ai_studio_model_id = "gemini-1.5-pro"
state.is_page_ready = True
state.page_instance = real_locks_mock_browser.page_instance
state.model_switching_lock = real_locks_mock_browser.model_switching_lock
state.params_cache_lock = real_locks_mock_browser.params_cache_lock
state.page_params_cache = {}
state.parsed_model_list = model_list
request = make_chat_request()
with patch("api_utils.server_state.state.logger", MagicMock()):
context = await initialize_request_context("req7", request)
assert context["parsed_model_list"] == model_list
@pytest.mark.asyncio
async def test_context_includes_all_required_fields(
self, real_locks_mock_browser, make_chat_request
):
"""Test that context includes all required fields."""
state.current_ai_studio_model_id = "gemini-1.5-pro"
state.is_page_ready = True
state.page_instance = real_locks_mock_browser.page_instance
state.model_switching_lock = real_locks_mock_browser.model_switching_lock
state.params_cache_lock = real_locks_mock_browser.params_cache_lock
state.page_params_cache = {}
state.parsed_model_list = []
request = make_chat_request()
with patch("api_utils.server_state.state.logger", MagicMock()):
context = await initialize_request_context("req8", request)
# Verify all required keys exist
required_keys = [
"logger",
"page",
"is_page_ready",
"parsed_model_list",
"current_ai_studio_model_id",
"model_switching_lock",
"page_params_cache",
"params_cache_lock",
"is_streaming",
"model_actually_switched",
"requested_model",
"model_id_to_use",
"needs_model_switching",
]
for key in required_keys:
assert key in context, f"Missing required key: {key}"
# Verify default flag values
assert context["model_actually_switched"] is False
assert context["model_id_to_use"] is None
assert context["needs_model_switching"] is False
@pytest.mark.asyncio
async def test_context_return_type_is_dict(
self, real_locks_mock_browser, make_chat_request
):
"""Test that context is returned as a dictionary."""
state.current_ai_studio_model_id = "gemini-1.5-pro"
state.is_page_ready = True
state.page_instance = real_locks_mock_browser.page_instance
state.model_switching_lock = real_locks_mock_browser.model_switching_lock
state.params_cache_lock = real_locks_mock_browser.params_cache_lock
state.page_params_cache = {}
state.parsed_model_list = []
request = make_chat_request()
with patch("api_utils.server_state.state.logger", MagicMock()):
context = await initialize_request_context("req9", request)
assert isinstance(context, dict)
@pytest.mark.asyncio
async def test_logger_message_format(
self, real_locks_mock_browser, make_chat_request
):
"""Test logger message format includes request details."""
state.current_ai_studio_model_id = "gemini-1.5-pro"
state.is_page_ready = True
state.page_instance = real_locks_mock_browser.page_instance
state.model_switching_lock = real_locks_mock_browser.model_switching_lock
state.params_cache_lock = real_locks_mock_browser.params_cache_lock
state.page_params_cache = {}
state.parsed_model_list = []
request = make_chat_request(model="test-model-123", stream=True)
with patch("api_utils.server_state.state.logger", MagicMock()) as mock_logger:
await initialize_request_context("test-req-abc", request)
# Verify log messages use debug, not info
log_calls = [call[0][0] for call in mock_logger.debug.call_args_list]
# Log should include [Request] tag and model/stream parameters
assert any("[Request]" in msg for msg in log_calls)
assert any("test-model-123" in msg for msg in log_calls)
assert any("Stream=True" in msg for msg in log_calls)
@pytest.mark.asyncio
async def test_locks_are_real_asyncio_locks(
self, real_locks_mock_browser, make_chat_request
):
"""Test that locks in context are real asyncio.Lock instances."""
state.current_ai_studio_model_id = "gemini-1.5-pro"
state.is_page_ready = True
state.page_instance = real_locks_mock_browser.page_instance
state.model_switching_lock = real_locks_mock_browser.model_switching_lock
state.params_cache_lock = real_locks_mock_browser.params_cache_lock
state.page_params_cache = {}
state.parsed_model_list = []
request = make_chat_request()
with patch("api_utils.server_state.state.logger", MagicMock()):
context = await initialize_request_context("req10", request)
# Verify locks are real asyncio.Lock instances
assert isinstance(context["model_switching_lock"], asyncio.Lock)
assert isinstance(context["params_cache_lock"], asyncio.Lock)
# Verify locks can be acquired
async with context["model_switching_lock"]:
assert context["model_switching_lock"].locked()
assert not context["model_switching_lock"].locked()
async with context["params_cache_lock"]:
assert context["params_cache_lock"].locked()
assert not context["params_cache_lock"].locked()
|