ronantakizawa/codereview-bench huggingface dataset:
Repository: home-assistant/core
File: tests/components/matter/common.py (Python)

Code before:
    client: MagicMock,
    event: EventType = EventType.ATTRIBUTE_UPDATED,
    data: Any = None,
) -> None:
    """Trigger a subscription callback."""
    # trigger callback on all subscribers
    for sub in client.subscribe_events.call_args_list:
        callback = sub.kwargs["callback"]
        event_filter = sub.kwargs.get("event_filter")
        if event_filter in (None, event):
            callback(event, data)
    await hass.async_block_till_done()


def snapshot_matter_entities(
    hass: HomeAssistant,
    entity_registry: er.EntityRegistry,
    snapshot: SnapshotAssertion,
    platform: Platform,
) -> None:
    """Snapshot Matter entities."""
    entities = hass.states.async_all(platform)
    for entity_state in entities:
        entity_entry = entity_registry.async_get(entity_state.entity_id)
        node_id = int(entity_entry.unique_id.split("-")[1], 16)
        fixture_name = FIXTURES[node_id - 1]
        assert entity_entry == snapshot(
            name=f"{fixture_name}-{entity_entry.entity_id}-entry"
        )
        assert entity_state == snapshot(
            name=f"{fixture_name}-{entity_entry.entity_id}-state"
        )


Review comment (refactor):
Extracting the fixture name to make it easier to identify snapshots as it is no longer provided by `pytest.parameterize`

Code after:
async def trigger_subscription_callback(
    hass: HomeAssistant,
    client: MagicMock,
    event: EventType = EventType.ATTRIBUTE_UPDATED,
    data: Any = None,
) -> None:
    """Trigger a subscription callback."""
    # trigger callback on all subscribers
    for sub in client.subscribe_events.call_args_list:
        callback = sub.kwargs["callback"]
        event_filter = sub.kwargs.get("event_filter")
        if event_filter in (None, event):
            callback(event, data)
    await hass.async_block_till_done()


@cache
def _get_fixture_name(node_id: int) -> dict[int, str]:
    """Get the fixture name for a given node ID."""
    for fixture_name in FIXTURES:
        fixture_data = load_and_parse_node_fixture(fixture_name)
        if fixture_data["node_id"] == node_id:
            return fixture_name

    raise KeyError(f"Fixture for node id {node_id} not found")


def snapshot_matter_entities(
    hass: HomeAssistant,
    entity_registry: er.EntityRegistry,
    snapshot: SnapshotAssertion,
    platform: Platform,
) -> None:
    """Snapshot Matter entities."""
    entities = hass.states.async_all(platform)
    for entity_state in entities:
        entity_entry = entity_registry.async_get(entity_state.entity_id)
        node_id = int(entity_entry.unique_id.split("-")[1], 16)
        fixture_name = _get_fixture_name(node_id)
        assert entity_entry == snapshot(
            name=f"{fixture_name}][{entity_entry.entity_id}-entry"
        )
        assert entity_state == snapshot(
            name=f"{fixture_name}][{entity_entry.entity_id}-state"
        )
Repository: home-assistant/core
File: homeassistant/components/mealie/todo.py (Python)

Code before:
                mutate_shopping_item.unit_id = item.unit_id

            await self.coordinator.client.update_shopping_item(
                mutate_shopping_item.item_id, mutate_shopping_item
            )

        await self.coordinator.async_refresh()

    @property
    def available(self) -> bool:
        """Return False if shopping list no longer available."""
        return super().available and self._shopping_list_id in self.coordinator.data

    async def async_get_shopping_list_items(self) -> ServiceResponse:
        """Get structured shopping list items."""
        client = self.coordinator.client
        try:
            shopping_items = await client.get_shopping_items(self._shopping_list_id)
        except MealieConnectionError as err:
            raise HomeAssistantError(
                translation_domain=DOMAIN,
                translation_key="connection_error",
            ) from err
        return {
            "name": self.shopping_list.name,
            "shopping_list_items": asdict(shopping_items),
        }


Review comment (suggestion):
`shopping_items` appears to have an `items` inside. Want to just flatten that one level? I think if there were other attributes on `shopping_items` they could also be merged into the top level.

Code after:
                mutate_shopping_item.unit_id = item.unit_id

            await self.coordinator.client.update_shopping_item(
                mutate_shopping_item.item_id, mutate_shopping_item
            )

        await self.coordinator.async_refresh()

    @property
    def available(self) -> bool:
        """Return False if shopping list no longer available."""
        return super().available and self._shopping_list_id in self.coordinator.data

    async def async_get_shopping_list_items(self) -> ServiceResponse:
        """Get structured shopping list items."""
        client = self.coordinator.client
        try:
            shopping_items = await client.get_shopping_items(self._shopping_list_id)
        except MealieConnectionError as err:
            raise HomeAssistantError(
                translation_domain=DOMAIN,
                translation_key="connection_error",
            ) from err
        return {
            "name": self.shopping_list.name,
            "items": [asdict(item) for item in shopping_items.items],
        }
Repository: home-assistant/core
File: homeassistant/components/mastodon/config_flow.py (Python)

Code before:
REAUTH_SCHEMA = vol.Schema(
    {
        vol.Required(
            CONF_ACCESS_TOKEN,
        ): TextSelector(TextSelectorConfig(type=TextSelectorType.PASSWORD)),
    }
)

EXAMPLE_URL = "https://mastodon.social"


def base_url_from_url(url: str) -> str:
    """Return the base url from a url."""
    return str(URL(url).origin())


class MastodonConfigFlow(ConfigFlow, domain=DOMAIN):
    """Handle a config flow."""

    VERSION = 1
    MINOR_VERSION = 2

    base_url: str | None = None
    client_id: str | None = None
    client_secret: str | None = None
    access_token: str | None = None

    def check_connection(
        self,
    ) -> tuple[
        InstanceV2 | Instance | None,
        Account | None,
        dict[str, str],
    ]:
        """Check connection to the Mastodon instance."""

        assert self.base_url is not None
        assert self.client_id is not None
        assert self.client_secret is not None
        assert self.access_token is not None

        try:
            client = create_mastodon_client(
                self.base_url,
                self.client_id,
                self.client_secret,
                self.access_token,
            )
            try:
                instance = client.instance_v2()
            except MastodonNotFoundError:

Review comment (security):
Do they need a default value?

```suggestion
    base_url: str
    client_id: str
    client_secret: str
    access_token: str
```

Code after:
REAUTH_SCHEMA = vol.Schema(
    {
        vol.Required(
            CONF_ACCESS_TOKEN,
        ): TextSelector(TextSelectorConfig(type=TextSelectorType.PASSWORD)),
    }
)

EXAMPLE_URL = "https://mastodon.social"


def base_url_from_url(url: str) -> str:
    """Return the base url from a url."""
    return str(URL(url).origin())


def remove_email_link(account_name: str) -> str:
    """Remove email link from account name."""

    # Replaces the @ with a HTML entity to prevent mailto links.
    return account_name.replace("@", "&#64;")


class MastodonConfigFlow(ConfigFlow, domain=DOMAIN):
    """Handle a config flow."""

    VERSION = 1
    MINOR_VERSION = 2

    base_url: str
    client_id: str
    client_secret: str
    access_token: str

    def check_connection(
        self,
    ) -> tuple[
        InstanceV2 | Instance | None,
        Account | None,
        dict[str, str],
    ]:
        """Check connection to the Mastodon instance."""
        try:
            client = create_mastodon_client(
                self.base_url,
                self.client_id,
                self.client_secret,
                self.access_token,
            )
            try:
                instance = client.instance_v2()
Repository: home-assistant/core
File: tests/components/smarla/test_init.py (Python)

Code before:
"""Test switch platform for Swing2Sleep Smarla integration."""

from unittest.mock import MagicMock

from pysmarlaapi.connection.exceptions import (
    AuthenticationException,
    ConnectionException,
)
import pytest

from homeassistant.config_entries import ConfigEntryState
from homeassistant.core import HomeAssistant

from . import setup_integration

from tests.common import MockConfigEntry


@pytest.mark.parametrize(
    ("exception", "expected_state"),
    [
        (AuthenticationException, ConfigEntryState.SETUP_ERROR),
        (ConnectionException, ConfigEntryState.SETUP_RETRY),
    ],
)
@pytest.mark.usefixtures("mock_federwiege")
async def test_init_exception(
    hass: HomeAssistant,
    mock_config_entry: MockConfigEntry,
    mock_refresh_token: MagicMock,
    exception: type[Exception],
    expected_state: ConfigEntryState,
) -> None:
    """Test init config setup exception."""
    mock_refresh_token.side_effect = exception

    assert not await setup_integration(hass, mock_config_entry)

    assert mock_config_entry.state is expected_state


async def test_init_auth_failure(
    hass: HomeAssistant,
    mock_config_entry: MockConfigEntry,
    mock_federwiege_cls: MagicMock,
) -> None:
    """Test behavior on invalid authentication during runtime."""
    assert await setup_integration(hass, mock_config_entry)

    assert mock_config_entry.state is ConfigEntryState.LOADED
    assert not any(mock_config_entry.async_get_active_flows(hass, {"reauth"}))

    auth_failed_cb = mock_federwiege_cls.call_args.args[2]
    await auth_failed_cb()
    await hass.async_block_till_done()
    assert any(mock_config_entry.async_get_active_flows(hass, {"reauth"}))


Review comment (security):
why are we not starting reauth during runtime? You could trigger that manually through `entry.async_start_reauth(hass)` if need be

Code after:
    mock_federwiege_cls: MagicMock,
) -> None:
    """Test behavior on invalid authentication during runtime."""
    invalid_auth_callback: Callable[[], Awaitable[None]] | None = None

    def mocked_federwiege(_1, _2, callback):
        nonlocal invalid_auth_callback
        invalid_auth_callback = callback
        return mock_federwiege_cls.return_value

    # Mock Federwiege class to gather authentication failure callback
    mock_federwiege_cls.side_effect = mocked_federwiege

    assert await setup_integration(hass, mock_config_entry)
    assert mock_config_entry.state is ConfigEntryState.LOADED

    # Check that config entry has no active reauth flows
    assert not any(mock_config_entry.async_get_active_flows(hass, {"reauth"}))

    # Simulate authentication failure during runtime
    assert invalid_auth_callback is not None
    await invalid_auth_callback()
    await hass.async_block_till_done()

    # Check that a reauth flow has been started
    assert any(mock_config_entry.async_get_active_flows(hass, {"reauth"}))
Repository: home-assistant/core
File: homeassistant/components/sony_projector/switch.py (Python)

Code before:

    def update(self) -> None:
        """Get the latest state from the projector."""
        try:
            self._attr_is_on = self._sdcp.get_power()
            self._attr_available = True
        except ConnectionRefusedError:
            _LOGGER.error("Projector connection refused")
            self._attr_available = False

    def turn_on(self, **kwargs: Any) -> None:
        """Turn the projector on."""
        _LOGGER.debug("Powering on projector '%s'", self.name)
        if self._sdcp.set_power(True):
            _LOGGER.debug("Powered on successfully")
            self._attr_is_on = True
        else:
            _LOGGER.error("Power on command was not successful")

    def turn_off(self, **kwargs: Any) -> None:
        """Turn the projector off."""
        _LOGGER.debug("Powering off projector '%s'", self.name)
        if self._sdcp.set_power(False):
            _LOGGER.debug("Powered off successfully")
            self._attr_is_on = False
        else:
            _LOGGER.error("Power off command was not successful")


Review comment (bug):
I'm pretty sure the previous behavior was wrong, return `STATE_ON`/`STATE_OFF` for `is_on` property.

Code after:
        self._attr_name = name

    def update(self) -> None:
        """Get the latest state from the projector."""
        try:
            self._attr_is_on = self._sdcp.get_power()
            self._attr_available = True
        except ConnectionRefusedError:
            _LOGGER.error("Projector connection refused")
            self._attr_available = False

    def turn_on(self, **kwargs: Any) -> None:
        """Turn the projector on."""
        _LOGGER.debug("Powering on projector '%s'", self.name)
        if self._sdcp.set_power(True):
            _LOGGER.debug("Powered on successfully")
            self._attr_is_on = True
        else:
            _LOGGER.error("Power on command was not successful")

    def turn_off(self, **kwargs: Any) -> None:
        """Turn the projector off."""
        _LOGGER.debug("Powering off projector '%s'", self.name)
        if self._sdcp.set_power(False):
            _LOGGER.debug("Powered off successfully")
            self._attr_is_on = False
        else:
            _LOGGER.error("Power off command was not successful")
Repository: home-assistant/core
File: homeassistant/components/hdfury/number.py (Python)

Code before:
async def async_setup_entry(
    hass: HomeAssistant,
    entry: HDFuryConfigEntry,
    async_add_entities: AddConfigEntryEntitiesCallback,
) -> None:
    """Set up numbers using the platform schema."""

    coordinator = entry.runtime_data

    async_add_entities(
        HDFuryNumber(coordinator, description)
        for description in NUMBERS
        if description.key in coordinator.data.config
    )


class HDFuryNumber(HDFuryEntity, NumberEntity):
    """Base HDFury Number Class."""

    entity_description: HDFuryNumberEntityDescription

    @property
    def native_value(self) -> float:
        """Set Number Value."""

        return float(self.coordinator.data.config[self.entity_description.key])

    async def async_set_native_value(self, value: float) -> None:
        """Set Number Value Event."""

        try:
            await self.entity_description.set_value_fn(
                self.coordinator.client, str(int(value))
            )
        except HDFuryError as error:
            raise HomeAssistantError(
                translation_domain=DOMAIN,
                translation_key="communication_error",
            ) from error

        await self.coordinator.async_request_refresh()


Review comment (suggestion):
If the data already is a float, you can use `cast(float, ...)` to be a tiny bit quicker

Code after:
async def async_setup_entry(
    hass: HomeAssistant,
    entry: HDFuryConfigEntry,
    async_add_entities: AddConfigEntryEntitiesCallback,
) -> None:
    """Set up numbers using the platform schema."""

    coordinator = entry.runtime_data

    async_add_entities(
        HDFuryNumber(coordinator, description)
        for description in NUMBERS
        if description.key in coordinator.data.config
    )


class HDFuryNumber(HDFuryEntity, NumberEntity):
    """Base HDFury Number Class."""

    entity_description: HDFuryNumberEntityDescription

    @property
    def native_value(self) -> float:
        """Return the current number value."""

        return float(self.coordinator.data.config[self.entity_description.key])

    async def async_set_native_value(self, value: float) -> None:
        """Set Number Value Event."""

        try:
            await self.entity_description.set_value_fn(
                self.coordinator.client, str(int(value))
            )
        except HDFuryError as error:
            raise HomeAssistantError(
                translation_domain=DOMAIN,
                translation_key="communication_error",
            ) from error

        await self.coordinator.async_request_refresh()
Repository: home-assistant/core
File: homeassistant/components/airos/button.py (Python)

Code before:
class AirOSRebootButton(AirOSEntity, ButtonEntity):
    """Button to reboot device."""

    _attr_has_entity_name = True
    _attr_device_class = ButtonDeviceClass.RESTART

    entity_description: ButtonEntityDescription

    def __init__(self, coordinator: AirOSDataUpdateCoordinator) -> None:
        """Initialize the AirOS client button."""
        super().__init__(coordinator)
        self.coordinator = coordinator
        self.entity_description = BUTTON_DESCRIPTION

        self._attr_unique_id = f"{coordinator.config_entry.unique_id}_reboot"

        self._attr_name = "Reboot Device"

    async def async_press(self) -> None:
        """Handle the button press to reboot the device."""
        result: bool = False
        try:
            await self.coordinator.airos_device.login()
            result = await self.coordinator.airos_device.reboot()

        except AirOSDataMissingError, AirOSDeviceConnectionError:
            _LOGGER.exception("Failed to reboot device")

        else:
            if not result:
                _LOGGER.error("Unable to reboot device")


Review comment (suggestion):
Since you're catching the exception, don't we want to raise an `HomeAssistantError` and notify the users about it? :)

Code after:

    entity_description: ButtonEntityDescription

    def __init__(
        self,
        coordinator: AirOSDataUpdateCoordinator,
        description: ButtonEntityDescription,
    ) -> None:
        """Initialize the AirOS client button."""
        super().__init__(coordinator)

        self.entity_description = description
        self._attr_unique_id = f"{coordinator.data.derived.mac}_{description.key}"

    async def async_press(self) -> None:
        """Handle the button press to reboot the device."""
        try:
            await self.coordinator.airos_device.login()
            result = await self.coordinator.airos_device.reboot()

        except AirOSException as err:
            raise HomeAssistantError(
                translation_domain=DOMAIN,
                translation_key="cannot_connect",
            ) from err

        if not result:
            raise HomeAssistantError(
                translation_domain=DOMAIN,
                translation_key="reboot_failed",
            ) from None
Repository: home-assistant/core
File: tests/components/airos/conftest.py (Python)

Code before:
    with patch(
        "homeassistant.components.airos.async_setup_entry", return_value=True
    ) as mock_setup_entry:
        yield mock_setup_entry


@pytest.fixture
def mock_airos_class() -> Generator[MagicMock]:
    """Fixture to mock the AirOS class itself."""
    with (
        patch("homeassistant.components.airos.AirOS8", autospec=True) as mock_class,
        patch("homeassistant.components.airos.config_flow.AirOS8", new=mock_class),
        patch("homeassistant.components.airos.coordinator.AirOS8", new=mock_class),
    ):
        yield mock_class


@pytest.fixture
def mock_airos_client(
    mock_airos_class: MagicMock, ap_fixture: AirOS8Data
) -> Generator[AsyncMock]:
    """Fixture to mock the AirOS API client."""
    client = mock_airos_class.return_value
    client.status = AsyncMock(return_value=ap_fixture)
    client.login = AsyncMock(return_value=True)
    client.reboot = AsyncMock(return_value=True)
    return client


@pytest.fixture
def mock_config_entry() -> MockConfigEntry:
    """Return the AirOS mocked config entry."""
    return MockConfigEntry(
        title="NanoStation",
        domain=DOMAIN,
        data={
            CONF_HOST: "1.1.1.1",
            CONF_PASSWORD: "test-password",
            CONF_USERNAME: "ubnt",
        },
        unique_id="01:23:45:67:89:AB",
    )


Review comment (suggestion):
```suggestion
    client.status.return_value = ap_fixture
    client.login.return_value = True
    client.reboot.return_value = True
```
Personally a big fan of the more inline notation

Code after:
    with patch(
        "homeassistant.components.airos.async_setup_entry", return_value=True
    ) as mock_setup_entry:
        yield mock_setup_entry


@pytest.fixture
def mock_airos_class() -> Generator[MagicMock]:
    """Fixture to mock the AirOS class itself."""
    with (
        patch("homeassistant.components.airos.AirOS8", autospec=True) as mock_class,
        patch("homeassistant.components.airos.config_flow.AirOS8", new=mock_class),
        patch("homeassistant.components.airos.coordinator.AirOS8", new=mock_class),
    ):
        yield mock_class


@pytest.fixture
def mock_airos_client(
    mock_airos_class: MagicMock, ap_fixture: AirOS8Data
) -> Generator[AsyncMock]:
    """Fixture to mock the AirOS API client."""
    client = mock_airos_class.return_value
    client.status.return_value = ap_fixture
    client.login.return_value = True
    client.reboot.return_value = True
    return client


@pytest.fixture
def mock_config_entry() -> MockConfigEntry:
    """Return the AirOS mocked config entry."""
    return MockConfigEntry(
        title="NanoStation",
        domain=DOMAIN,
        data={
            CONF_HOST: "1.1.1.1",
            CONF_PASSWORD: "test-password",
            CONF_USERNAME: "ubnt",
        },
        unique_id="01:23:45:67:89:AB",
    )
Repository: home-assistant/core
File: homeassistant/components/proxmoxve/diagnostics.py (Python)

Code before:
        device_registry, config_entry.entry_id
    )

    for device in registry_devices:
        entities = []

        registry_entities = er.async_entries_for_device(
            entity_registry,
            device_id=device.id,
            include_disabled_entities=True,
        )

        for entity_entry in registry_entities:
            state_dict = None
            if state := hass.states.get(entity_entry.entity_id):
                state_dict = dict(state.as_dict())
                state_dict.pop("context", None)

            entities.append({"entry": asdict(entity_entry), "state": state_dict})

        devices.append({"device": asdict(device), "entities": entities})

    return {
        "config_entry": async_redact_data(config_entry.as_dict(), TO_REDACT),
        "devices": devices,
    }


Review comment (bug):
I am wondering what the benefit of entities and devices is in the diagnostic information. When something is wrong it's always easy to ask the user for information about these things, but a platform like this is ideal to just get the coordinator data and know what the integration based the states off of instead of just all the entities

Code after:
"""Diagnostics support for Proxmox VE."""

from __future__ import annotations

from dataclasses import asdict
from typing import Any

from homeassistant.components.diagnostics import async_redact_data
from homeassistant.const import CONF_HOST, CONF_PASSWORD, CONF_USERNAME
from homeassistant.core import HomeAssistant

from . import ProxmoxConfigEntry

TO_REDACT = [CONF_USERNAME, CONF_PASSWORD, CONF_HOST]


async def async_get_config_entry_diagnostics(
    hass: HomeAssistant, config_entry: ProxmoxConfigEntry
) -> dict[str, Any]:
    """Return diagnostics for a Proxmox VE config entry."""

    return {
        "config_entry": async_redact_data(config_entry.as_dict(), TO_REDACT),
        "devices": {
            node: asdict(node_data)
            for node, node_data in config_entry.runtime_data.data.items()
        },
    }