mirror of
https://github.com/home-assistant/core.git
synced 2026-09-30 22:23:25 +01:00
376 lines
12 KiB
Python
376 lines
12 KiB
Python
"""Tests for the Anthropic integration."""
|
|
|
|
from pathlib import Path
|
|
import re
|
|
from unittest.mock import AsyncMock, patch
|
|
|
|
from anthropic.types import Message, TextBlock, Usage
|
|
from freezegun import freeze_time
|
|
import probatio
|
|
import pytest
|
|
from syrupy.assertion import SnapshotAssertion
|
|
|
|
from homeassistant.components import ai_task, media_source
|
|
from homeassistant.core import HomeAssistant
|
|
from homeassistant.exceptions import HomeAssistantError
|
|
from homeassistant.helpers import entity_registry as er, selector
|
|
|
|
from . import create_content_block
|
|
|
|
from tests.common import MockConfigEntry
|
|
|
|
|
|
@pytest.mark.usefixtures("mock_init_component")
|
|
async def test_generate_data(
|
|
hass: HomeAssistant,
|
|
mock_config_entry: MockConfigEntry,
|
|
mock_create_stream: AsyncMock,
|
|
entity_registry: er.EntityRegistry,
|
|
) -> None:
|
|
"""Test AI Task data generation."""
|
|
entity_id = "ai_task.claude_ai_task"
|
|
|
|
# Ensure entity is linked to the subentry
|
|
entity_entry = entity_registry.async_get(entity_id)
|
|
ai_task_entry = next(
|
|
iter(
|
|
entry
|
|
for entry in mock_config_entry.subentries.values()
|
|
if entry.subentry_type == "ai_task_data"
|
|
)
|
|
)
|
|
assert entity_entry is not None
|
|
assert entity_entry.config_entry_id == mock_config_entry.entry_id
|
|
assert entity_entry.config_subentry_id == ai_task_entry.subentry_id
|
|
|
|
mock_create_stream.return_value = [create_content_block(0, ["The test data"])]
|
|
|
|
result = await ai_task.async_generate_data(
|
|
hass,
|
|
task_name="Test Task",
|
|
entity_id=entity_id,
|
|
instructions="Generate test data",
|
|
)
|
|
|
|
assert result.data == "The test data"
|
|
|
|
|
|
@pytest.mark.usefixtures("mock_init_component")
|
|
async def test_translation_key(
|
|
entity_registry: er.EntityRegistry,
|
|
) -> None:
|
|
"""Test entity translation key."""
|
|
entry = entity_registry.async_get("ai_task.claude_ai_task")
|
|
assert entry is not None
|
|
assert entry.translation_key == "ai_task_data"
|
|
|
|
|
|
@pytest.mark.usefixtures("mock_init_component")
|
|
async def test_empty_data(
|
|
hass: HomeAssistant,
|
|
mock_create_stream: AsyncMock,
|
|
) -> None:
|
|
"""Test AI Task data generation but the data returned is empty."""
|
|
mock_create_stream.return_value = [create_content_block(0, [""])]
|
|
|
|
with pytest.raises(
|
|
HomeAssistantError, match="Last content in chat log is not an AssistantContent"
|
|
):
|
|
await ai_task.async_generate_data(
|
|
hass,
|
|
task_name="Test Task",
|
|
entity_id="ai_task.claude_ai_task",
|
|
instructions="Generate test data",
|
|
)
|
|
|
|
|
|
@pytest.mark.usefixtures("mock_init_component")
|
|
async def test_stream_wrong_type(
|
|
hass: HomeAssistant,
|
|
mock_create_stream: AsyncMock,
|
|
) -> None:
|
|
"""Test error if the response is not a stream."""
|
|
mock_create_stream.return_value = Message(
|
|
type="message",
|
|
id="message_id",
|
|
model="claude-fable-5",
|
|
role="assistant",
|
|
content=[TextBlock(type="text", text="This is not a stream")],
|
|
usage=Usage(input_tokens=42, output_tokens=42),
|
|
)
|
|
|
|
with pytest.raises(HomeAssistantError, match="Expected a stream of messages"):
|
|
await ai_task.async_generate_data(
|
|
hass,
|
|
task_name="Test Task",
|
|
entity_id="ai_task.claude_ai_task",
|
|
instructions="Generate test data",
|
|
)
|
|
|
|
|
|
@pytest.mark.usefixtures("mock_init_component")
|
|
async def test_generate_invalid_structured_data(
|
|
hass: HomeAssistant,
|
|
mock_create_stream: AsyncMock,
|
|
) -> None:
|
|
"""Test AI Task with an invalid JSON response."""
|
|
mock_create_stream.return_value = [
|
|
create_content_block(0, ["INVALID JSON RESPONSE"])
|
|
]
|
|
|
|
with pytest.raises(
|
|
HomeAssistantError, match="Error with Claude structured response"
|
|
):
|
|
await ai_task.async_generate_data(
|
|
hass,
|
|
task_name="Test Task",
|
|
entity_id="ai_task.claude_ai_task",
|
|
instructions="Generate test data",
|
|
structure=probatio.Schema(
|
|
{
|
|
probatio.Required("characters"): selector.selector(
|
|
{
|
|
"text": {
|
|
"multiple": True,
|
|
}
|
|
}
|
|
)
|
|
},
|
|
),
|
|
)
|
|
|
|
|
|
@freeze_time("2026-01-01 12:00:00")
|
|
@pytest.mark.usefixtures("mock_init_component")
|
|
async def test_generate_structured_data(
|
|
hass: HomeAssistant,
|
|
mock_create_stream: AsyncMock,
|
|
snapshot: SnapshotAssertion,
|
|
) -> None:
|
|
"""Test AI Task structured data generation."""
|
|
mock_create_stream.return_value = [
|
|
create_content_block(0, ['{"charac', 'ters": ["Mario', '", "Luigi"]}'])
|
|
]
|
|
|
|
result = await ai_task.async_generate_data(
|
|
hass,
|
|
task_name="Test Task",
|
|
entity_id="ai_task.claude_ai_task",
|
|
instructions="Generate test data",
|
|
structure=probatio.Schema(
|
|
{
|
|
probatio.Required("characters"): selector.selector(
|
|
{
|
|
"text": {
|
|
"multiple": True,
|
|
}
|
|
}
|
|
)
|
|
},
|
|
),
|
|
)
|
|
|
|
assert result.data == {"characters": ["Mario", "Luigi"]}
|
|
assert mock_create_stream.call_args.kwargs.copy() == snapshot
|
|
|
|
|
|
@pytest.mark.usefixtures("mock_init_component")
|
|
async def test_generate_data_with_attachments(
|
|
hass: HomeAssistant,
|
|
mock_create_stream: AsyncMock,
|
|
) -> None:
|
|
"""Test AI Task data generation with attachments."""
|
|
entity_id = "ai_task.claude_ai_task"
|
|
|
|
mock_create_stream.return_value = [create_content_block(0, ["Hi there!"])]
|
|
|
|
# Test with attachments
|
|
with (
|
|
patch(
|
|
"homeassistant.components.media_source.async_resolve_media",
|
|
side_effect=[
|
|
media_source.PlayMedia(
|
|
url="http://example.com/doorbell_snapshot.jpg",
|
|
mime_type="image/jpg",
|
|
path=Path("doorbell_snapshot.jpg"),
|
|
),
|
|
media_source.PlayMedia(
|
|
url="http://example.com/context.pdf",
|
|
mime_type="application/pdf",
|
|
path=Path("context.pdf"),
|
|
),
|
|
],
|
|
),
|
|
patch("pathlib.Path.exists", return_value=True),
|
|
patch("pathlib.Path.read_bytes", return_value=b"fake_image_data"),
|
|
):
|
|
result = await ai_task.async_generate_data(
|
|
hass,
|
|
task_name="Test Task",
|
|
entity_id=entity_id,
|
|
instructions="Test prompt",
|
|
attachments=[
|
|
{"media_content_id": "media-source://media/doorbell_snapshot.jpg"},
|
|
{"media_content_id": "media-source://media/context.pdf"},
|
|
],
|
|
)
|
|
|
|
assert result.data == "Hi there!"
|
|
|
|
# Verify that the create stream was called with the correct parameters
|
|
# The last call should have the user message with attachments
|
|
call_args = mock_create_stream.call_args
|
|
assert call_args is not None
|
|
|
|
# Check that the input includes the attachments
|
|
input_messages = call_args[1]["messages"]
|
|
assert len(input_messages) > 0
|
|
|
|
# Find the user message with attachments
|
|
user_message_with_attachments = input_messages[-2]
|
|
|
|
assert user_message_with_attachments is not None
|
|
assert isinstance(user_message_with_attachments["content"], list)
|
|
assert len(user_message_with_attachments["content"]) == 3 # Text + attachments
|
|
|
|
text_block, image_block, document_block = user_message_with_attachments["content"]
|
|
|
|
# Text block
|
|
assert text_block["type"] == "text"
|
|
assert text_block["text"] == "Test prompt"
|
|
|
|
# Image attachment
|
|
assert image_block["type"] == "image"
|
|
assert image_block["source"] == {
|
|
"data": "ZmFrZV9pbWFnZV9kYXRh",
|
|
"media_type": "image/jpeg",
|
|
"type": "base64",
|
|
}
|
|
|
|
# Document attachment (ignore extra metadata like cache_control)
|
|
assert document_block["type"] == "document"
|
|
assert document_block["source"]["data"] == "ZmFrZV9pbWFnZV9kYXRh"
|
|
assert document_block["source"]["media_type"] == "application/pdf"
|
|
assert document_block["source"]["type"] == "base64"
|
|
|
|
|
|
@pytest.mark.usefixtures("mock_init_component")
|
|
async def test_generate_data_invalid_attachments(
|
|
hass: HomeAssistant,
|
|
mock_create_stream: AsyncMock,
|
|
) -> None:
|
|
"""Test AI Task data generation with attachments of unsupported type."""
|
|
entity_id = "ai_task.claude_ai_task"
|
|
|
|
mock_create_stream.return_value = [create_content_block(0, ["Hi there!"])]
|
|
|
|
# Test path that doesn't exist
|
|
with (
|
|
patch(
|
|
"homeassistant.components.media_source.async_resolve_media",
|
|
side_effect=[
|
|
media_source.PlayMedia(
|
|
url="http://example.com/doorbell_snapshot.jpg",
|
|
mime_type="image/jpeg",
|
|
path=Path("doorbell_snapshot.jpg"),
|
|
)
|
|
],
|
|
),
|
|
patch("pathlib.Path.exists", return_value=False),
|
|
pytest.raises(
|
|
HomeAssistantError, match="`doorbell_snapshot.jpg` does not exist"
|
|
),
|
|
):
|
|
await ai_task.async_generate_data(
|
|
hass,
|
|
task_name="Test Task",
|
|
entity_id=entity_id,
|
|
instructions="Test prompt",
|
|
attachments=[
|
|
{"media_content_id": "media-source://media/doorbell_snapshot.jpg"},
|
|
],
|
|
)
|
|
|
|
# Test unsupported file type
|
|
with (
|
|
patch(
|
|
"homeassistant.components.media_source.async_resolve_media",
|
|
side_effect=[
|
|
media_source.PlayMedia(
|
|
url="http://example.com/doorbell_snapshot.txt",
|
|
mime_type=None,
|
|
path=Path("doorbell_snapshot.txt"),
|
|
)
|
|
],
|
|
),
|
|
patch("pathlib.Path.exists", return_value=True),
|
|
patch(
|
|
"homeassistant.components.anthropic.entity.guess_file_type",
|
|
return_value=("text/plain", None),
|
|
),
|
|
pytest.raises(
|
|
HomeAssistantError,
|
|
match=re.escape(
|
|
"The Claude Haiku 4.5 model does not support"
|
|
" text/plain file types"
|
|
" (for `doorbell_snapshot.txt`)"
|
|
),
|
|
),
|
|
):
|
|
await ai_task.async_generate_data(
|
|
hass,
|
|
task_name="Test Task",
|
|
entity_id=entity_id,
|
|
instructions="Test prompt",
|
|
attachments=[
|
|
{"media_content_id": "media-source://media/doorbell_snapshot.txt"},
|
|
],
|
|
)
|
|
|
|
|
|
@pytest.mark.usefixtures("mock_init_component")
|
|
async def test_generate_data_with_attachments_whitespace_instructions(
|
|
hass: HomeAssistant,
|
|
mock_create_stream: AsyncMock,
|
|
) -> None:
|
|
"""Test whitespace-only instructions with attachments produce no text block.
|
|
|
|
The API rejects whitespace-only text blocks, so the user message should
|
|
contain only the attachment.
|
|
"""
|
|
entity_id = "ai_task.claude_ai_task"
|
|
|
|
mock_create_stream.return_value = [create_content_block(0, ["Hi there!"])]
|
|
|
|
with (
|
|
patch(
|
|
"homeassistant.components.media_source.async_resolve_media",
|
|
side_effect=[
|
|
media_source.PlayMedia(
|
|
url="http://example.com/doorbell_snapshot.jpg",
|
|
mime_type="image/jpg",
|
|
path=Path("doorbell_snapshot.jpg"),
|
|
),
|
|
],
|
|
),
|
|
patch("pathlib.Path.exists", return_value=True),
|
|
patch("pathlib.Path.read_bytes", return_value=b"fake_image_data"),
|
|
):
|
|
result = await ai_task.async_generate_data(
|
|
hass,
|
|
task_name="Test Task",
|
|
entity_id=entity_id,
|
|
instructions=" ",
|
|
attachments=[
|
|
{"media_content_id": "media-source://media/doorbell_snapshot.jpg"},
|
|
],
|
|
)
|
|
|
|
assert result.data == "Hi there!"
|
|
|
|
input_messages = mock_create_stream.call_args[1]["messages"]
|
|
user_message = input_messages[-2]
|
|
assert user_message["role"] == "user"
|
|
assert isinstance(user_message["content"], list)
|
|
assert [block["type"] for block in user_message["content"]] == ["image"]
|