tuya-local/tests/test_config_flow.py

"""Tests for the config flow."""

import pytest
import voluptuous as vol
from homeassistant.const import CONF_HOST, CONF_NAME
from homeassistant.data_entry_flow import FlowResultType
from pytest_homeassistant_custom_component.common import MockConfigEntry

from custom_components.tuya_local import (
    async_migrate_entry,
    config_flow,
    get_device_unique_id,
)
from custom_components.tuya_local.const import (
    CONF_DEVICE_CID,
    CONF_DEVICE_ID,
    CONF_LOCAL_KEY,
    CONF_POLL_ONLY,
    CONF_PROTOCOL_VERSION,
    CONF_TYPE,
    DOMAIN,
)

# Designed to contain "special" characters that users constantly suspect.
TESTKEY = ")<jO<@)'P1|kR$Kd"


@pytest.fixture(autouse=True)
def auto_enable_custom_integrations(enable_custom_integrations):
    yield


@pytest.fixture(autouse=True)
def prevent_task_creation(mocker):
    mocker.patch("custom_components.tuya_local.device.TuyaLocalDevice.register_entity")
    yield


@pytest.fixture
def bypass_setup(mocker):
    """Prevent actual setup of the integration after config flow."""
    mocker.patch("custom_components.tuya_local.async_setup_entry", return_value=True)
    yield


@pytest.fixture
def bypass_data_fetch(mocker):
    """Prevent actual data fetching from the device."""
    mocker.patch("tinytuya.Device.status", return_value={"1": True})
    yield


@pytest.mark.asyncio
async def test_init_entry(hass, bypass_data_fetch):
    """Test initialisation of the config flow."""
    entry = MockConfigEntry(
        domain=DOMAIN,
        version=11,
        title="test",
        data={
            CONF_DEVICE_ID: "deviceid",
            CONF_HOST: "hostname",
            CONF_LOCAL_KEY: TESTKEY,
            CONF_POLL_ONLY: False,
            CONF_PROTOCOL_VERSION: "auto",
            CONF_TYPE: "kogan_kahtp_heater",
            CONF_DEVICE_CID: None,
        },
        options={},
    )
    entry.add_to_hass(hass)
    await hass.config_entries.async_setup(entry.entry_id)
    await hass.async_block_till_done()
    assert hass.states.get("climate.test")
    assert hass.states.get("lock.test_child_lock")


@pytest.mark.asyncio
async def test_migrate_entry(hass, mocker):
    """Test migration from old entry format."""
    mock_device = mocker.MagicMock()
    mock_device.async_inferred_type = mocker.AsyncMock(
        return_value="goldair_gpph_heater"
    )
    mocker.patch("custom_components.tuya_local.setup_device", return_value=mock_device)

    entry = MockConfigEntry(
        domain=DOMAIN,
        version=1,
        title="test",
        data={
            CONF_DEVICE_ID: "deviceid",
            CONF_HOST: "hostname",
            CONF_LOCAL_KEY: TESTKEY,
            CONF_TYPE: "auto",
            "climate": True,
            "child_lock": True,
            "display_light": True,
        },
    )
    entry.add_to_hass(hass)
    assert await async_migrate_entry(hass, entry)

    mock_device.async_inferred_type = mocker.AsyncMock(return_value=None)
    mock_device.reset_mock()

    entry = MockConfigEntry(
        domain=DOMAIN,
        version=1,
        title="test2",
        data={
            CONF_DEVICE_ID: "deviceid",
            CONF_HOST: "hostname",
            CONF_LOCAL_KEY: TESTKEY,
            CONF_TYPE: "unknown",
            "climate": False,
        },
    )
    entry.add_to_hass(hass)
    assert not await async_migrate_entry(hass, entry)
    mock_device.reset_mock()

    entry = MockConfigEntry(
        domain=DOMAIN,
        version=2,
        title="test3",
        data={
            CONF_DEVICE_ID: "deviceid",
            CONF_HOST: "hostname",
            CONF_LOCAL_KEY: TESTKEY,
            CONF_TYPE: "auto",
        },
        options={
            "climate": False,
        },
    )
    entry.add_to_hass(hass)
    assert not await async_migrate_entry(hass, entry)

    mock_device.async_inferred_type = mocker.AsyncMock(return_value="smartplugv1")
    mock_device.reset_mock()

    entry = MockConfigEntry(
        domain=DOMAIN,
        version=3,
        title="test4",
        data={
            CONF_DEVICE_ID: "deviceid",
            CONF_HOST: "hostname",
            CONF_LOCAL_KEY: TESTKEY,
            CONF_TYPE: "smartplugv1",
        },
        options={
            "switch": True,
        },
    )
    entry.add_to_hass(hass)
    assert await async_migrate_entry(hass, entry)

    mock_device.async_inferred_type = mocker.AsyncMock(return_value="smartplugv2")
    mock_device.reset_mock()

    entry = MockConfigEntry(
        domain=DOMAIN,
        version=3,
        title="test5",
        data={
            CONF_DEVICE_ID: "deviceid",
            CONF_HOST: "hostname",
            CONF_LOCAL_KEY: TESTKEY,
            CONF_TYPE: "smartplugv1",
        },
        options={
            "switch": True,
        },
    )
    entry.add_to_hass(hass)
    assert await async_migrate_entry(hass, entry)

    mock_device.async_inferred_type = mocker.AsyncMock(
        return_value="goldair_dehumidifier"
    )
    mock_device.reset_mock()

    entry = MockConfigEntry(
        domain=DOMAIN,
        version=4,
        title="test6",
        data={
            CONF_DEVICE_ID: "deviceid",
            CONF_HOST: "hostname",
            CONF_LOCAL_KEY: TESTKEY,
            CONF_TYPE: "goldair_dehumidifier",
        },
        options={
            "humidifier": True,
            "fan": True,
            "light": True,
            "lock": False,
            "switch": True,
        },
    )
    entry.add_to_hass(hass)
    assert await async_migrate_entry(hass, entry)

    mock_device.async_inferred_type = mocker.AsyncMock(
        return_value="grid_connect_usb_double_power_point"
    )
    mock_device.reset_mock()

    entry = MockConfigEntry(
        domain=DOMAIN,
        version=6,
        title="test7",
        data={
            CONF_DEVICE_ID: "deviceid",
            CONF_HOST: "hostname",
            CONF_LOCAL_KEY: TESTKEY,
            CONF_TYPE: "grid_connect_usb_double_power_point",
        },
        options={
            "switch_main_switch": True,
            "switch_left_outlet": True,
            "switch_right_outlet": True,
        },
    )
    entry.add_to_hass(hass)
    assert await async_migrate_entry(hass, entry)


@pytest.mark.asyncio
async def test_flow_user_init(hass, mocker):
    """Test the initialisation of the form in the first page of the manual config flow path."""
    result = await hass.config_entries.flow.async_init(
        DOMAIN, context={"source": "local"}
    )
    expected = {
        "data_schema": mocker.ANY,
        "description_placeholders": mocker.ANY,
        "errors": {},
        "flow_id": mocker.ANY,
        "handler": DOMAIN,
        "step_id": "local",
        "type": "form",
        "last_step": mocker.ANY,
        "preview": mocker.ANY,
    }
    assert expected == result
    # Check the schema.  Simple comparison does not work since they are not
    # the same object
    try:
        result["data_schema"](
            {CONF_DEVICE_ID: "test", CONF_LOCAL_KEY: TESTKEY, CONF_HOST: "test"}
        )
    except vol.MultipleInvalid:
        assert False
    try:
        result["data_schema"]({CONF_DEVICE_ID: "missing_some"})
        assert False
    except vol.MultipleInvalid:
        pass


@pytest.mark.asyncio
async def test_async_test_connection_valid(hass, mocker):
    """Test that device is returned when connection is valid."""
    mock_device = mocker.patch(
        "custom_components.tuya_local.config_flow.TuyaLocalDevice"
    )
    mock_instance = mocker.AsyncMock()
    mock_instance.has_returned_state = True
    mock_instance.pause = mocker.MagicMock()
    mock_instance.resume = mocker.MagicMock()
    mock_device.return_value = mock_instance
    hass.data[DOMAIN] = {"deviceid": {"device": mock_instance}}

    device = await config_flow.async_test_connection(
        {
            CONF_DEVICE_ID: "deviceid",
            CONF_LOCAL_KEY: TESTKEY,
            CONF_HOST: "hostname",
            CONF_PROTOCOL_VERSION: "auto",
        },
        hass,
    )
    assert device == mock_instance
    mock_instance.pause.assert_called_once()
    mock_instance.resume.assert_called_once()


@pytest.mark.asyncio
async def test_async_test_connection_for_subdevice_valid(hass, mocker):
    """Test that subdevice is returned when connection is valid."""
    mock_device = mocker.patch(
        "custom_components.tuya_local.config_flow.TuyaLocalDevice"
    )
    mock_instance = mocker.AsyncMock()
    mock_instance.has_returned_state = True
    mock_instance.pause = mocker.MagicMock()
    mock_instance.resume = mocker.MagicMock()
    mock_device.return_value = mock_instance
    hass.data[DOMAIN] = {"subdeviceid": {"device": mock_instance}}

    device = await config_flow.async_test_connection(
        {
            CONF_DEVICE_ID: "deviceid",
            CONF_LOCAL_KEY: TESTKEY,
            CONF_HOST: "hostname",
            CONF_PROTOCOL_VERSION: "auto",
            CONF_DEVICE_CID: "subdeviceid",
        },
        hass,
    )
    assert device == mock_instance
    mock_instance.pause.assert_called_once()
    mock_instance.resume.assert_called_once()


@pytest.mark.asyncio
async def test_async_test_connection_invalid(hass, mocker):
    """Test that None is returned when connection is invalid."""
    mock_device = mocker.patch(
        "custom_components.tuya_local.config_flow.TuyaLocalDevice"
    )
    mock_instance = mocker.AsyncMock()
    mock_instance.has_returned_state = False
    mock_instance._api = mocker.MagicMock()
    mock_device.return_value = mock_instance
    device = await config_flow.async_test_connection(
        {
            CONF_DEVICE_ID: "deviceid",
            CONF_LOCAL_KEY: TESTKEY,
            CONF_HOST: "hostname",
            CONF_PROTOCOL_VERSION: "auto",
        },
        hass,
    )
    assert device is None


@pytest.mark.asyncio
async def test_flow_user_init_invalid_config(hass, mocker):
    """Test errors populated when config is invalid."""
    mocker.patch(
        "custom_components.tuya_local.config_flow.async_test_connection",
        return_value=None,
    )
    flow = await hass.config_entries.flow.async_init(
        DOMAIN, context={"source": "local"}
    )
    result = await hass.config_entries.flow.async_configure(
        flow["flow_id"],
        user_input={
            CONF_DEVICE_ID: "deviceid",
            CONF_HOST: "hostname",
            CONF_LOCAL_KEY: "badkey",
            CONF_PROTOCOL_VERSION: "auto",
            CONF_POLL_ONLY: False,
        },
    )
    assert {"base": "connection"} == result["errors"]


def setup_device_mock(mock, mocker, failure=False, type="test"):
    mock_type = mocker.MagicMock()
    mock_type.legacy_type = type
    mock_type.config_type = type
    mock_type.match_quality.return_value = 100
    mock_type.product_display_entries.return_value = [(None, None)]
    mock.async_possible_types = mocker.AsyncMock(
        return_value=[mock_type] if not failure else []
    )


@pytest.mark.asyncio
async def test_flow_user_init_data_valid(hass, mocker):
    """Test we advance to the next step when connection config is valid."""
    mock_device = mocker.MagicMock()
    mock_device._protocol_configured = "auto"
    setup_device_mock(mock_device, mocker)
    mocker.patch(
        "custom_components.tuya_local.config_flow.async_test_connection",
        return_value=mock_device,
    )

    flow = await hass.config_entries.flow.async_init(
        DOMAIN, context={"source": "local"}
    )
    result = await hass.config_entries.flow.async_configure(
        flow["flow_id"],
        user_input={
            CONF_DEVICE_ID: "deviceid",
            CONF_HOST: "hostname",
            CONF_LOCAL_KEY: TESTKEY,
        },
    )
    assert "form" == result["type"]
    assert "select_type" == result["step_id"]


@pytest.mark.asyncio
async def test_flow_select_type_init(hass, mocker):
    """Test the initialisation of the form in the 2nd step of the config flow."""
    mock_device = mocker.patch.object(config_flow.ConfigFlowHandler, "device")

    setup_device_mock(mock_device, mocker)

    result = await hass.config_entries.flow.async_init(
        DOMAIN, context={"source": "select_type"}
    )
    expected = {
        "data_schema": mocker.ANY,
        "description_placeholders": None,
        "errors": None,
        "flow_id": mocker.ANY,
        "handler": DOMAIN,
        "step_id": "select_type",
        "type": "form",
        "last_step": mocker.ANY,
        "preview": mocker.ANY,
    }
    assert expected == result
    # Check the schema.  Simple comparison does not work since they are not
    # the same object
    try:
        result["data_schema"]({CONF_TYPE: "test||||"})
    except vol.MultipleInvalid:
        assert False
    try:
        result["data_schema"]({CONF_TYPE: "not_test||||"})
        assert False
    except vol.MultipleInvalid:
        pass


@pytest.mark.asyncio
async def test_flow_select_type_aborts_when_no_match(hass, mocker):
    """Test the flow aborts when an unsupported device is used."""
    mock_device = mocker.patch.object(config_flow.ConfigFlowHandler, "device")
    setup_device_mock(mock_device, mocker, failure=True)

    result = await hass.config_entries.flow.async_init(
        DOMAIN, context={"source": "select_type"}
    )

    assert result["type"] == "abort"
    assert result["reason"] == "not_supported"


@pytest.mark.asyncio
async def test_flow_select_type_data_valid(hass, mocker):
    """Test the flow continues when valid data is supplied."""
    mock_device = mocker.patch.object(config_flow.ConfigFlowHandler, "device")

    setup_device_mock(mock_device, mocker, type="smartplugv1")

    flow = await hass.config_entries.flow.async_init(
        DOMAIN, context={"source": "select_type"}
    )
    result = await hass.config_entries.flow.async_configure(
        flow["flow_id"],
        user_input={CONF_TYPE: "smartplugv1||||"},
    )
    assert "form" == result["type"]
    assert "choose_entities" == result["step_id"]


@pytest.mark.asyncio
async def test_flow_choose_entities_init(hass, mocker):
    """Test the initialisation of the form in the 3rd step of the config flow."""

    mocker.patch.dict(config_flow.ConfigFlowHandler.data, {CONF_TYPE: "smartplugv1"})
    result = await hass.config_entries.flow.async_init(
        DOMAIN, context={"source": "choose_entities"}
    )

    expected = {
        "data_schema": mocker.ANY,
        "description_placeholders": None,
        "errors": None,
        "flow_id": mocker.ANY,
        "handler": DOMAIN,
        "step_id": "choose_entities",
        "type": "form",
        "last_step": mocker.ANY,
        "preview": mocker.ANY,
    }
    assert expected == result
    # Check the schema.  Simple comparison does not work since they are not
    # the same object
    try:
        result["data_schema"]({CONF_NAME: "test"})
    except vol.MultipleInvalid:
        assert False
    try:
        result["data_schema"]({"climate": True})
        assert False
    except vol.MultipleInvalid:
        pass


@pytest.mark.asyncio
async def test_flow_choose_entities_creates_config_entry(hass, bypass_setup, mocker):
    """Test the flow ends when data is valid."""

    mocker.patch.dict(
        config_flow.ConfigFlowHandler.data,
        {
            CONF_DEVICE_ID: "deviceid",
            CONF_LOCAL_KEY: TESTKEY,
            CONF_HOST: "hostname",
            CONF_POLL_ONLY: False,
            CONF_PROTOCOL_VERSION: "auto",
            CONF_TYPE: "kogan_kahtp_heater",
            CONF_DEVICE_CID: None,
        },
    )
    flow = await hass.config_entries.flow.async_init(
        DOMAIN, context={"source": "choose_entities"}
    )
    result = await hass.config_entries.flow.async_configure(
        flow["flow_id"],
        user_input={
            CONF_NAME: "test",
        },
    )
    expected = {
        "version": 13,
        "minor_version": mocker.ANY,
        "context": {"source": "choose_entities"},
        "type": FlowResultType.CREATE_ENTRY,
        "flow_id": mocker.ANY,
        "handler": DOMAIN,
        "title": "test",
        "description": None,
        "description_placeholders": None,
        "result": mocker.ANY,
        "subentries": (),
        "options": {},
        "data": {
            CONF_DEVICE_ID: "deviceid",
            CONF_HOST: "hostname",
            CONF_LOCAL_KEY: TESTKEY,
            CONF_POLL_ONLY: False,
            CONF_PROTOCOL_VERSION: "auto",
            CONF_TYPE: "kogan_kahtp_heater",
            CONF_DEVICE_CID: None,
        },
    }
    assert expected == result


@pytest.mark.asyncio
async def test_options_flow_init(hass, bypass_data_fetch):
    """Test config flow options."""
    config_entry = MockConfigEntry(
        domain=DOMAIN,
        version=13,
        unique_id="uniqueid",
        data={
            CONF_DEVICE_ID: "deviceid",
            CONF_HOST: "hostname",
            CONF_LOCAL_KEY: TESTKEY,
            CONF_NAME: "test",
            CONF_POLL_ONLY: False,
            CONF_PROTOCOL_VERSION: "auto",
            CONF_TYPE: "smartplugv1",
            CONF_DEVICE_CID: "",
        },
    )
    config_entry.add_to_hass(hass)

    assert await hass.config_entries.async_setup(config_entry.entry_id)
    await hass.async_block_till_done()

    # show initial form
    result = await hass.config_entries.options.async_init(config_entry.entry_id)
    assert "form" == result["type"]
    assert "user" == result["step_id"]
    assert {} == result["errors"]
    assert result["data_schema"](
        {
            CONF_HOST: "hostname",
            CONF_LOCAL_KEY: TESTKEY,
        }
    )


@pytest.mark.asyncio
async def test_options_flow_modifies_config(hass, bypass_setup, mocker):
    mock_device = mocker.MagicMock()
    mocker.patch(
        "custom_components.tuya_local.config_flow.async_test_connection",
        return_value=mock_device,
    )

    config_entry = MockConfigEntry(
        domain=DOMAIN,
        version=13,
        unique_id="uniqueid",
        data={
            CONF_DEVICE_ID: "deviceid",
            CONF_HOST: "hostname",
            CONF_LOCAL_KEY: TESTKEY,
            CONF_NAME: "test",
            CONF_POLL_ONLY: False,
            CONF_PROTOCOL_VERSION: "auto",
            CONF_TYPE: "ble_pt216_temp_humidity",
            CONF_DEVICE_CID: "subdeviceid",
        },
    )
    config_entry.add_to_hass(hass)

    assert await hass.config_entries.async_setup(config_entry.entry_id)
    await hass.async_block_till_done()
    # show initial form
    form = await hass.config_entries.options.async_init(config_entry.entry_id)
    # submit updated config
    result = await hass.config_entries.options.async_configure(
        form["flow_id"],
        user_input={
            CONF_HOST: "new_hostname",
            CONF_LOCAL_KEY: "new_key",
            CONF_POLL_ONLY: False,
            CONF_PROTOCOL_VERSION: 3.3,
        },
    )
    expected = {
        CONF_HOST: "new_hostname",
        CONF_LOCAL_KEY: "new_key",
        CONF_POLL_ONLY: False,
        CONF_PROTOCOL_VERSION: 3.3,
    }
    assert "create_entry" == result["type"]
    assert "" == result["title"]
    assert expected == result["data"]


@pytest.mark.asyncio
async def test_options_flow_fails_when_connection_fails(
    hass, bypass_data_fetch, mocker
):
    mocker.patch(
        "custom_components.tuya_local.config_flow.async_test_connection",
        return_value=None,
    )
    config_entry = MockConfigEntry(
        domain=DOMAIN,
        version=13,
        unique_id="uniqueid",
        data={
            CONF_DEVICE_ID: "deviceid",
            CONF_HOST: "hostname",
            CONF_LOCAL_KEY: TESTKEY,
            CONF_NAME: "test",
            CONF_POLL_ONLY: False,
            CONF_PROTOCOL_VERSION: "auto",
            CONF_TYPE: "smartplugv1",
            CONF_DEVICE_CID: "",
        },
    )
    config_entry.add_to_hass(hass)

    assert await hass.config_entries.async_setup(config_entry.entry_id)
    await hass.async_block_till_done()
    # show initial form
    form = await hass.config_entries.options.async_init(config_entry.entry_id)
    # submit updated config
    result = await hass.config_entries.options.async_configure(
        form["flow_id"],
        user_input={
            CONF_HOST: "new_hostname",
            CONF_LOCAL_KEY: "new_key",
        },
    )
    assert "form" == result["type"]
    assert "user" == result["step_id"]
    assert {"base": "connection"} == result["errors"]


@pytest.mark.asyncio
async def test_options_flow_fails_when_config_is_missing(hass, mocker):
    mock_device = mocker.MagicMock()
    mocker.patch(
        "custom_components.tuya_local.config_flow.async_test_connection",
        return_value=mock_device,
    )

    config_entry = MockConfigEntry(
        domain=DOMAIN,
        version=13,
        unique_id="uniqueid",
        data={
            CONF_DEVICE_ID: "deviceid",
            CONF_HOST: "hostname",
            CONF_LOCAL_KEY: TESTKEY,
            CONF_NAME: "test",
            CONF_POLL_ONLY: False,
            CONF_PROTOCOL_VERSION: "auto",
            CONF_TYPE: "non_existing",
        },
    )
    config_entry.add_to_hass(hass)

    await hass.config_entries.async_setup(config_entry.entry_id)
    await hass.async_block_till_done()
    # show initial form
    result = await hass.config_entries.options.async_init(config_entry.entry_id)
    assert result["type"] == "abort"
    assert result["reason"] == "not_supported"


def test_migration_gets_correct_device_id():
    """Test that migration gets the correct device id."""
    # Normal device
    entry = MockConfigEntry(
        domain=DOMAIN,
        version=1,
        title="test",
        data={
            CONF_DEVICE_ID: "deviceid",
            CONF_HOST: "hostname",
            CONF_LOCAL_KEY: TESTKEY,
            CONF_TYPE: "auto",
        },
    )
    assert get_device_unique_id(entry) == "deviceid"