Move Google Gen AI fixture to allow reuse (#146921)

This commit is contained in:
Paulus Schoutsen 2025-06-15 23:00:27 -04:00 committed by GitHub
parent fa21269f0d
commit 8498928e47
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 21 additions and 21 deletions

View File

@ -1,6 +1,7 @@
"""Tests helpers.""" """Tests helpers."""
from unittest.mock import Mock, patch from collections.abc import Generator
from unittest.mock import AsyncMock, Mock, patch
import pytest import pytest
@ -77,3 +78,22 @@ async def mock_init_component(
async def setup_ha(hass: HomeAssistant) -> None: async def setup_ha(hass: HomeAssistant) -> None:
"""Set up Home Assistant.""" """Set up Home Assistant."""
assert await async_setup_component(hass, "homeassistant", {}) assert await async_setup_component(hass, "homeassistant", {})
@pytest.fixture
def mock_send_message_stream() -> Generator[AsyncMock]:
"""Mock stream response."""
async def mock_generator(stream):
for value in stream:
yield value
with patch(
"google.genai.chats.AsyncChat.send_message_stream",
AsyncMock(),
) as mock_send_message_stream:
mock_send_message_stream.side_effect = lambda **kwargs: mock_generator(
mock_send_message_stream.return_value.pop(0)
)
yield mock_send_message_stream

View File

@ -1,6 +1,5 @@
"""Tests for the Google Generative AI Conversation integration conversation platform.""" """Tests for the Google Generative AI Conversation integration conversation platform."""
from collections.abc import Generator
from unittest.mock import AsyncMock, patch from unittest.mock import AsyncMock, patch
from freezegun import freeze_time from freezegun import freeze_time
@ -41,25 +40,6 @@ def mock_ulid_tools():
yield yield
@pytest.fixture
def mock_send_message_stream() -> Generator[AsyncMock]:
"""Mock stream response."""
async def mock_generator(stream):
for value in stream:
yield value
with patch(
"google.genai.chats.AsyncChat.send_message_stream",
AsyncMock(),
) as mock_send_message_stream:
mock_send_message_stream.side_effect = lambda **kwargs: mock_generator(
mock_send_message_stream.return_value.pop(0)
)
yield mock_send_message_stream
@pytest.mark.parametrize( @pytest.mark.parametrize(
("error"), ("error"),
[ [