"""The tests for the MQTT switch platform."""
import json
from unittest.mock import ANY

from asynctest import patch
import pytest

from homeassistant.components import mqtt, switch
from homeassistant.components.mqtt.discovery import async_start
from homeassistant.const import (
    ATTR_ASSUMED_STATE,
    STATE_OFF,
    STATE_ON,
    STATE_UNAVAILABLE,
)
import homeassistant.core as ha
from homeassistant.setup import async_setup_component

from tests.common import (
    MockConfigEntry,
    async_fire_mqtt_message,
    async_mock_mqtt_component,
    mock_coro,
    mock_registry,
)
from tests.components.switch import common


@pytest.fixture
def mock_publish(hass):
    """Initialize components."""
    yield hass.loop.run_until_complete(async_mock_mqtt_component(hass))


async def test_controlling_state_via_topic(hass, mock_publish):
    """Test the controlling state via topic."""
    assert await async_setup_component(
        hass,
        switch.DOMAIN,
        {
            switch.DOMAIN: {
                "platform": "mqtt",
                "name": "test",
                "state_topic": "state-topic",
                "command_topic": "command-topic",
                "payload_on": 1,
                "payload_off": 0,
            }
        },
    )

    state = hass.states.get("switch.test")
    assert state.state == STATE_OFF
    assert not state.attributes.get(ATTR_ASSUMED_STATE)

    async_fire_mqtt_message(hass, "state-topic", "1")

    state = hass.states.get("switch.test")
    assert state.state == STATE_ON

    async_fire_mqtt_message(hass, "state-topic", "0")

    state = hass.states.get("switch.test")
    assert state.state == STATE_OFF


async def test_sending_mqtt_commands_and_optimistic(hass, mock_publish):
    """Test the sending MQTT commands in optimistic mode."""
    fake_state = ha.State("switch.test", "on")

    with patch(
        "homeassistant.helpers.restore_state.RestoreEntity" ".async_get_last_state",
        return_value=mock_coro(fake_state),
    ):
        assert await async_setup_component(
            hass,
            switch.DOMAIN,
            {
                switch.DOMAIN: {
                    "platform": "mqtt",
                    "name": "test",
                    "command_topic": "command-topic",
                    "payload_on": "beer on",
                    "payload_off": "beer off",
                    "qos": "2",
                }
            },
        )

    state = hass.states.get("switch.test")
    assert state.state == STATE_ON
    assert state.attributes.get(ATTR_ASSUMED_STATE)

    await common.async_turn_on(hass, "switch.test")

    mock_publish.async_publish.assert_called_once_with(
        "command-topic", "beer on", 2, False
    )
    mock_publish.async_publish.reset_mock()
    state = hass.states.get("switch.test")
    assert state.state == STATE_ON

    await common.async_turn_off(hass, "switch.test")

    mock_publish.async_publish.assert_called_once_with(
        "command-topic", "beer off", 2, False
    )
    state = hass.states.get("switch.test")
    assert state.state == STATE_OFF


async def test_controlling_state_via_topic_and_json_message(hass, mock_publish):
    """Test the controlling state via topic and JSON message."""
    assert await async_setup_component(
        hass,
        switch.DOMAIN,
        {
            switch.DOMAIN: {
                "platform": "mqtt",