邹永赫
45 lines
1.4 KiB
Python
45 lines
1.4 KiB
Python
from unittest.mock import patch
|
|
|
|
import pytest
|
|
|
|
from astrbot.core.provider.sources.gemini_source import ProviderGoogleGenAI
|
|
|
|
|
|
def _build_provider(provider_settings: dict) -> ProviderGoogleGenAI:
|
|
with patch.object(ProviderGoogleGenAI, "_init_client", lambda self: None):
|
|
return ProviderGoogleGenAI(
|
|
{
|
|
"key": ["test-key"],
|
|
"model": "gemini-3-pro-image-preview",
|
|
"timeout": 30,
|
|
"gm_safety_settings": {},
|
|
},
|
|
provider_settings,
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_prepare_query_config_keeps_image_modalities_for_non_stream_requests():
|
|
provider = _build_provider({"streaming_response": True})
|
|
|
|
config = await provider._prepare_query_config(
|
|
payloads={"messages": [], "model": "gemini-3-pro-image-preview"},
|
|
modalities=["TEXT", "IMAGE"],
|
|
streaming=False,
|
|
)
|
|
|
|
assert config.response_modalities == ["TEXT", "IMAGE"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_prepare_query_config_downgrades_image_modalities_for_stream_requests():
|
|
provider = _build_provider({"streaming_response": False})
|
|
|
|
config = await provider._prepare_query_config(
|
|
payloads={"messages": [], "model": "gemini-3-pro-image-preview"},
|
|
modalities=["TEXT", "IMAGE"],
|
|
streaming=True,
|
|
)
|
|
|
|
assert config.response_modalities == ["TEXT"]
|