From 97edbaa85f97384ca47b5029092f1fee00c8e164 Mon Sep 17 00:00:00 2001
From: Erik Montnemery <erik@montnemery.com>
Date: Thu, 10 Dec 2020 21:30:07 +0100
Subject: [PATCH] Small cleanup of MQTT (#44110)

* Use relative imports of mqtt component in mqtt platforms

* Correct parameters to _async_setup_entity

* Lint
---
 .../components/mqtt/alarm_control_panel.py    |  2 +-
 .../components/mqtt/binary_sensor.py          |  3 ++-
 homeassistant/components/mqtt/camera.py       |  3 ++-
 homeassistant/components/mqtt/climate.py      |  3 ++-
 homeassistant/components/mqtt/cover.py        |  3 ++-
 .../components/mqtt/device_automation.py      |  2 +-
 .../mqtt/device_tracker/schema_discovery.py   |  3 ++-
 .../mqtt/device_tracker/schema_yaml.py        |  2 +-
 .../components/mqtt/device_trigger.py         |  2 +-
 homeassistant/components/mqtt/discovery.py    |  2 +-
 homeassistant/components/mqtt/fan.py          |  3 ++-
 .../components/mqtt/light/__init__.py         |  8 ++-----
 .../components/mqtt/light/schema_basic.py     | 24 +++++++++----------
 .../components/mqtt/light/schema_json.py      | 24 +++++++++----------
 .../components/mqtt/light/schema_template.py  | 24 +++++++++----------
 homeassistant/components/mqtt/lock.py         |  3 ++-
 homeassistant/components/mqtt/scene.py        |  3 ++-
 homeassistant/components/mqtt/sensor.py       |  3 ++-
 homeassistant/components/mqtt/subscription.py |  2 +-
 homeassistant/components/mqtt/switch.py       |  5 ++--
 homeassistant/components/mqtt/tag.py          |  2 +-
 homeassistant/components/mqtt/trigger.py      |  3 ++-
 .../components/mqtt/vacuum/__init__.py        | 12 ++++------
 .../components/mqtt/vacuum/schema_legacy.py   | 16 ++++++-------
 .../components/mqtt/vacuum/schema_state.py    | 24 +++++++++----------
 25 files changed, 92 insertions(+), 89 deletions(-)

diff --git a/homeassistant/components/mqtt/alarm_control_panel.py b/homeassistant/components/mqtt/alarm_control_panel.py
index 35d0e1fb42e..edf383a6819 100644
--- a/homeassistant/components/mqtt/alarm_control_panel.py
+++ b/homeassistant/components/mqtt/alarm_control_panel.py
@@ -4,7 +4,6 @@ import re
 
 import voluptuous as vol
 
-from homeassistant.components import mqtt
 import homeassistant.components.alarm_control_panel as alarm
 from homeassistant.components.alarm_control_panel.const import (
     SUPPORT_ALARM_ARM_AWAY,
@@ -48,6 +47,7 @@ from . import (
     MqttEntityDeviceInfo,
     subscription,
 )
+from .. import mqtt
 from .debug_info import log_messages
 from .discovery import MQTT_DISCOVERY_NEW, clear_discovery_hash
 
diff --git a/homeassistant/components/mqtt/binary_sensor.py b/homeassistant/components/mqtt/binary_sensor.py
index 339b41a9ddc..e081423d590 100644
--- a/homeassistant/components/mqtt/binary_sensor.py
+++ b/homeassistant/components/mqtt/binary_sensor.py
@@ -4,7 +4,7 @@ import logging
 
 import voluptuous as vol
 
-from homeassistant.components import binary_sensor, mqtt
+from homeassistant.components import binary_sensor
 from homeassistant.components.binary_sensor import (
     DEVICE_CLASSES_SCHEMA,
     BinarySensorEntity,
@@ -40,6 +40,7 @@ from . import (
     MqttEntityDeviceInfo,
     subscription,
 )
+from .. import mqtt
 from .debug_info import log_messages
 from .discovery import MQTT_DISCOVERY_NEW, clear_discovery_hash
 
diff --git a/homeassistant/components/mqtt/camera.py b/homeassistant/components/mqtt/camera.py
index 82e5cb8b272..e8783f74bd4 100644
--- a/homeassistant/components/mqtt/camera.py
+++ b/homeassistant/components/mqtt/camera.py
@@ -3,7 +3,7 @@ import logging
 
 import voluptuous as vol
 
-from homeassistant.components import camera, mqtt
+from homeassistant.components import camera
 from homeassistant.components.camera import Camera
 from homeassistant.const import CONF_DEVICE, CONF_NAME, CONF_UNIQUE_ID
 from homeassistant.core import callback
@@ -23,6 +23,7 @@ from . import (
     MqttEntityDeviceInfo,
     subscription,
 )
+from .. import mqtt
 from .debug_info import log_messages
 from .discovery import MQTT_DISCOVERY_NEW, clear_discovery_hash
 
diff --git a/homeassistant/components/mqtt/climate.py b/homeassistant/components/mqtt/climate.py
index 8b762a82f02..c5835f8e7c7 100644
--- a/homeassistant/components/mqtt/climate.py
+++ b/homeassistant/components/mqtt/climate.py
@@ -3,7 +3,7 @@ import logging
 
 import voluptuous as vol
 
-from homeassistant.components import climate, mqtt
+from homeassistant.components import climate
 from homeassistant.components.climate import (
     PLATFORM_SCHEMA as CLIMATE_PLATFORM_SCHEMA,
     ClimateEntity,
@@ -64,6 +64,7 @@ from . import (
     MqttEntityDeviceInfo,
     subscription,
 )
+from .. import mqtt
 from .debug_info import log_messages
 from .discovery import MQTT_DISCOVERY_NEW, clear_discovery_hash
 
diff --git a/homeassistant/components/mqtt/cover.py b/homeassistant/components/mqtt/cover.py
index c3a78133246..25fcf0ad0d2 100644
--- a/homeassistant/components/mqtt/cover.py
+++ b/homeassistant/components/mqtt/cover.py
@@ -3,7 +3,7 @@ import logging
 
 import voluptuous as vol
 
-from homeassistant.components import cover, mqtt
+from homeassistant.components import cover
 from homeassistant.components.cover import (
     ATTR_POSITION,
     ATTR_TILT_POSITION,
@@ -51,6 +51,7 @@ from . import (
     MqttEntityDeviceInfo,
     subscription,
 )
+from .. import mqtt
 from .debug_info import log_messages
 from .discovery import MQTT_DISCOVERY_NEW, clear_discovery_hash
 
diff --git a/homeassistant/components/mqtt/device_automation.py b/homeassistant/components/mqtt/device_automation.py
index 4fcfd8f66f2..c064cca599d 100644
--- a/homeassistant/components/mqtt/device_automation.py
+++ b/homeassistant/components/mqtt/device_automation.py
@@ -3,11 +3,11 @@ import logging
 
 import voluptuous as vol
 
-from homeassistant.components import mqtt
 from homeassistant.helpers.device_registry import EVENT_DEVICE_REGISTRY_UPDATED
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
 
 from . import ATTR_DISCOVERY_HASH, device_trigger
+from .. import mqtt
 from .discovery import MQTT_DISCOVERY_NEW, clear_discovery_hash
 
 _LOGGER = logging.getLogger(__name__)
diff --git a/homeassistant/components/mqtt/device_tracker/schema_discovery.py b/homeassistant/components/mqtt/device_tracker/schema_discovery.py
index aa45f8f92b2..4de2ae4fa6d 100644
--- a/homeassistant/components/mqtt/device_tracker/schema_discovery.py
+++ b/homeassistant/components/mqtt/device_tracker/schema_discovery.py
@@ -3,7 +3,7 @@ import logging
 
 import voluptuous as vol
 
-from homeassistant.components import device_tracker, mqtt
+from homeassistant.components import device_tracker
 from homeassistant.components.device_tracker import SOURCE_TYPES
 from homeassistant.components.device_tracker.config_entry import TrackerEntity
 from homeassistant.const import (
@@ -29,6 +29,7 @@ from .. import (
     MqttEntityDeviceInfo,
     subscription,
 )
+from ... import mqtt
 from ..const import ATTR_DISCOVERY_HASH, CONF_QOS, CONF_STATE_TOPIC
 from ..debug_info import log_messages
 from ..discovery import MQTT_DISCOVERY_NEW, clear_discovery_hash
diff --git a/homeassistant/components/mqtt/device_tracker/schema_yaml.py b/homeassistant/components/mqtt/device_tracker/schema_yaml.py
index 520bced2385..f871ac89c2d 100644
--- a/homeassistant/components/mqtt/device_tracker/schema_yaml.py
+++ b/homeassistant/components/mqtt/device_tracker/schema_yaml.py
@@ -2,12 +2,12 @@
 
 import voluptuous as vol
 
-from homeassistant.components import mqtt
 from homeassistant.components.device_tracker import PLATFORM_SCHEMA, SOURCE_TYPES
 from homeassistant.const import CONF_DEVICES, STATE_HOME, STATE_NOT_HOME
 from homeassistant.core import callback
 import homeassistant.helpers.config_validation as cv
 
+from ... import mqtt
 from ..const import CONF_QOS
 
 CONF_PAYLOAD_HOME = "payload_home"
diff --git a/homeassistant/components/mqtt/device_trigger.py b/homeassistant/components/mqtt/device_trigger.py
index 676252c3134..9fa51bebf09 100644
--- a/homeassistant/components/mqtt/device_trigger.py
+++ b/homeassistant/components/mqtt/device_trigger.py
@@ -5,7 +5,6 @@ from typing import Callable, List, Optional
 import attr
 import voluptuous as vol
 
-from homeassistant.components import mqtt
 from homeassistant.components.automation import AutomationActionType
 from homeassistant.components.device_automation import TRIGGER_BASE_SCHEMA
 from homeassistant.const import CONF_DEVICE_ID, CONF_DOMAIN, CONF_PLATFORM, CONF_TYPE
@@ -28,6 +27,7 @@ from . import (
     debug_info,
     trigger as mqtt_trigger,
 )
+from .. import mqtt
 from .discovery import MQTT_DISCOVERY_UPDATED, clear_discovery_hash
 
 _LOGGER = logging.getLogger(__name__)
diff --git a/homeassistant/components/mqtt/discovery.py b/homeassistant/components/mqtt/discovery.py
index ca576f83d2a..5452d15aa30 100644
--- a/homeassistant/components/mqtt/discovery.py
+++ b/homeassistant/components/mqtt/discovery.py
@@ -6,12 +6,12 @@ import logging
 import re
 import time
 
-from homeassistant.components import mqtt
 from homeassistant.const import CONF_DEVICE, CONF_PLATFORM
 from homeassistant.helpers.dispatcher import async_dispatcher_send
 from homeassistant.helpers.typing import HomeAssistantType
 from homeassistant.loader import async_get_mqtt
 
+from .. import mqtt
 from .abbreviations import ABBREVIATIONS, DEVICE_ABBREVIATIONS
 from .const import (
     ATTR_DISCOVERY_HASH,
diff --git a/homeassistant/components/mqtt/fan.py b/homeassistant/components/mqtt/fan.py
index 14469e415e0..96d5fe720c3 100644
--- a/homeassistant/components/mqtt/fan.py
+++ b/homeassistant/components/mqtt/fan.py
@@ -3,7 +3,7 @@ import logging
 
 import voluptuous as vol
 
-from homeassistant.components import fan, mqtt
+from homeassistant.components import fan
 from homeassistant.components.fan import (
     ATTR_SPEED,
     SPEED_HIGH,
@@ -43,6 +43,7 @@ from . import (
     MqttEntityDeviceInfo,
     subscription,
 )
+from .. import mqtt
 from .debug_info import log_messages
 from .discovery import MQTT_DISCOVERY_NEW, clear_discovery_hash
 
diff --git a/homeassistant/components/mqtt/light/__init__.py b/homeassistant/components/mqtt/light/__init__.py
index 2375fb86e5a..393cb2fcf13 100644
--- a/homeassistant/components/mqtt/light/__init__.py
+++ b/homeassistant/components/mqtt/light/__init__.py
@@ -4,16 +4,12 @@ import logging
 import voluptuous as vol
 
 from homeassistant.components import light
-from homeassistant.components.mqtt import ATTR_DISCOVERY_HASH
-from homeassistant.components.mqtt.discovery import (
-    MQTT_DISCOVERY_NEW,
-    clear_discovery_hash,
-)
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.reload import async_setup_reload_service
 from homeassistant.helpers.typing import ConfigType, HomeAssistantType
 
-from .. import DOMAIN, PLATFORMS
+from .. import ATTR_DISCOVERY_HASH, DOMAIN, PLATFORMS
+from ..discovery import MQTT_DISCOVERY_NEW, clear_discovery_hash
 from .schema import CONF_SCHEMA, MQTT_LIGHT_SCHEMA_SCHEMA
 from .schema_basic import PLATFORM_SCHEMA_BASIC, async_setup_entity_basic
 from .schema_json import PLATFORM_SCHEMA_JSON, async_setup_entity_json
diff --git a/homeassistant/components/mqtt/light/schema_basic.py b/homeassistant/components/mqtt/light/schema_basic.py
index 4796652f57e..00ad2671391 100644
--- a/homeassistant/components/mqtt/light/schema_basic.py
+++ b/homeassistant/components/mqtt/light/schema_basic.py
@@ -3,7 +3,6 @@ import logging
 
 import voluptuous as vol
 
-from homeassistant.components import mqtt
 from homeassistant.components.light import (
     ATTR_BRIGHTNESS,
     ATTR_COLOR_TEMP,
@@ -17,17 +16,6 @@ from homeassistant.components.light import (
     SUPPORT_WHITE_VALUE,
     LightEntity,
 )
-from homeassistant.components.mqtt import (
-    CONF_COMMAND_TOPIC,
-    CONF_QOS,
-    CONF_RETAIN,
-    CONF_STATE_TOPIC,
-    MqttAttributes,
-    MqttAvailability,
-    MqttDiscoveryUpdate,
-    MqttEntityDeviceInfo,
-    subscription,
-)
 from homeassistant.const import (
     CONF_DEVICE,
     CONF_NAME,
@@ -43,6 +31,18 @@ import homeassistant.helpers.config_validation as cv
 from homeassistant.helpers.restore_state import RestoreEntity
 import homeassistant.util.color as color_util
 
+from .. import (
+    CONF_COMMAND_TOPIC,
+    CONF_QOS,
+    CONF_RETAIN,
+    CONF_STATE_TOPIC,
+    MqttAttributes,
+    MqttAvailability,
+    MqttDiscoveryUpdate,
+    MqttEntityDeviceInfo,
+    subscription,
+)
+from ... import mqtt
 from ..debug_info import log_messages
 from .schema import MQTT_LIGHT_SCHEMA_SCHEMA
 
diff --git a/homeassistant/components/mqtt/light/schema_json.py b/homeassistant/components/mqtt/light/schema_json.py
index bba5605348b..bb10fd52ae7 100644
--- a/homeassistant/components/mqtt/light/schema_json.py
+++ b/homeassistant/components/mqtt/light/schema_json.py
@@ -4,7 +4,6 @@ import logging
 
 import voluptuous as vol
 
-from homeassistant.components import mqtt
 from homeassistant.components.light import (
     ATTR_BRIGHTNESS,
     ATTR_COLOR_TEMP,
@@ -24,17 +23,6 @@ from homeassistant.components.light import (
     SUPPORT_WHITE_VALUE,
     LightEntity,
 )
-from homeassistant.components.mqtt import (
-    CONF_COMMAND_TOPIC,
-    CONF_QOS,
-    CONF_RETAIN,
-    CONF_STATE_TOPIC,
-    MqttAttributes,
-    MqttAvailability,
-    MqttDiscoveryUpdate,
-    MqttEntityDeviceInfo,
-    subscription,
-)
 from homeassistant.const import (
     CONF_BRIGHTNESS,
     CONF_COLOR_TEMP,
@@ -54,6 +42,18 @@ from homeassistant.helpers.restore_state import RestoreEntity
 from homeassistant.helpers.typing import ConfigType
 import homeassistant.util.color as color_util
 
+from .. import (
+    CONF_COMMAND_TOPIC,
+    CONF_QOS,
+    CONF_RETAIN,
+    CONF_STATE_TOPIC,
+    MqttAttributes,
+    MqttAvailability,
+    MqttDiscoveryUpdate,
+    MqttEntityDeviceInfo,
+    subscription,
+)
+from ... import mqtt
 from ..debug_info import log_messages
 from .schema import MQTT_LIGHT_SCHEMA_SCHEMA
 from .schema_basic import CONF_BRIGHTNESS_SCALE
diff --git a/homeassistant/components/mqtt/light/schema_template.py b/homeassistant/components/mqtt/light/schema_template.py
index faf987881b9..e6b22da5af0 100644
--- a/homeassistant/components/mqtt/light/schema_template.py
+++ b/homeassistant/components/mqtt/light/schema_template.py
@@ -3,7 +3,6 @@ import logging
 
 import voluptuous as vol
 
-from homeassistant.components import mqtt
 from homeassistant.components.light import (
     ATTR_BRIGHTNESS,
     ATTR_COLOR_TEMP,
@@ -21,17 +20,6 @@ from homeassistant.components.light import (
     SUPPORT_WHITE_VALUE,
     LightEntity,
 )
-from homeassistant.components.mqtt import (
-    CONF_COMMAND_TOPIC,
-    CONF_QOS,
-    CONF_RETAIN,
-    CONF_STATE_TOPIC,
-    MqttAttributes,
-    MqttAvailability,
-    MqttDiscoveryUpdate,
-    MqttEntityDeviceInfo,
-    subscription,
-)
 from homeassistant.const import (
     CONF_DEVICE,
     CONF_NAME,
@@ -45,6 +33,18 @@ import homeassistant.helpers.config_validation as cv
 from homeassistant.helpers.restore_state import RestoreEntity
 import homeassistant.util.color as color_util
 
+from .. import (
+    CONF_COMMAND_TOPIC,
+    CONF_QOS,
+    CONF_RETAIN,
+    CONF_STATE_TOPIC,
+    MqttAttributes,
+    MqttAvailability,
+    MqttDiscoveryUpdate,
+    MqttEntityDeviceInfo,
+    subscription,
+)
+from ... import mqtt
 from ..debug_info import log_messages
 from .schema import MQTT_LIGHT_SCHEMA_SCHEMA
 
diff --git a/homeassistant/components/mqtt/lock.py b/homeassistant/components/mqtt/lock.py
index aea1e40b0f9..712f2e0e376 100644
--- a/homeassistant/components/mqtt/lock.py
+++ b/homeassistant/components/mqtt/lock.py
@@ -3,7 +3,7 @@ import logging
 
 import voluptuous as vol
 
-from homeassistant.components import lock, mqtt
+from homeassistant.components import lock
 from homeassistant.components.lock import LockEntity
 from homeassistant.const import (
     CONF_DEVICE,
@@ -32,6 +32,7 @@ from . import (
     MqttEntityDeviceInfo,
     subscription,
 )
+from .. import mqtt
 from .debug_info import log_messages
 from .discovery import MQTT_DISCOVERY_NEW, clear_discovery_hash
 
diff --git a/homeassistant/components/mqtt/scene.py b/homeassistant/components/mqtt/scene.py
index 4f4380332fd..673eb169b19 100644
--- a/homeassistant/components/mqtt/scene.py
+++ b/homeassistant/components/mqtt/scene.py
@@ -3,7 +3,7 @@ import logging
 
 import voluptuous as vol
 
-from homeassistant.components import mqtt, scene
+from homeassistant.components import scene
 from homeassistant.components.scene import Scene
 from homeassistant.const import CONF_ICON, CONF_NAME, CONF_PAYLOAD_ON, CONF_UNIQUE_ID
 import homeassistant.helpers.config_validation as cv
@@ -21,6 +21,7 @@ from . import (
     MqttAvailability,
     MqttDiscoveryUpdate,
 )
+from .. import mqtt
 from .discovery import MQTT_DISCOVERY_NEW, clear_discovery_hash
 
 _LOGGER = logging.getLogger(__name__)
diff --git a/homeassistant/components/mqtt/sensor.py b/homeassistant/components/mqtt/sensor.py
index ffd34cef8c9..1fda8986ef7 100644
--- a/homeassistant/components/mqtt/sensor.py
+++ b/homeassistant/components/mqtt/sensor.py
@@ -5,7 +5,7 @@ from typing import Optional
 
 import voluptuous as vol
 
-from homeassistant.components import mqtt, sensor
+from homeassistant.components import sensor
 from homeassistant.components.sensor import DEVICE_CLASSES_SCHEMA
 from homeassistant.const import (
     CONF_DEVICE,
@@ -38,6 +38,7 @@ from . import (
     MqttEntityDeviceInfo,
     subscription,
 )
+from .. import mqtt
 from .debug_info import log_messages
 from .discovery import MQTT_DISCOVERY_NEW, clear_discovery_hash
 
diff --git a/homeassistant/components/mqtt/subscription.py b/homeassistant/components/mqtt/subscription.py
index 24c1c6ff3a1..c61c30c922e 100644
--- a/homeassistant/components/mqtt/subscription.py
+++ b/homeassistant/components/mqtt/subscription.py
@@ -4,11 +4,11 @@ from typing import Any, Callable, Dict, Optional
 
 import attr
 
-from homeassistant.components import mqtt
 from homeassistant.helpers.typing import HomeAssistantType
 from homeassistant.loader import bind_hass
 
 from . import debug_info
+from .. import mqtt
 from .const import DEFAULT_QOS
 from .models import MessageCallbackType
 
diff --git a/homeassistant/components/mqtt/switch.py b/homeassistant/components/mqtt/switch.py
index 761f19ef054..76019680110 100644
--- a/homeassistant/components/mqtt/switch.py
+++ b/homeassistant/components/mqtt/switch.py
@@ -3,7 +3,7 @@ import logging
 
 import voluptuous as vol
 
-from homeassistant.components import mqtt, switch
+from homeassistant.components import switch
 from homeassistant.components.switch import SwitchEntity
 from homeassistant.const import (
     CONF_DEVICE,
@@ -37,6 +37,7 @@ from . import (
     MqttEntityDeviceInfo,
     subscription,
 )
+from .. import mqtt
 from .debug_info import log_messages
 from .discovery import MQTT_DISCOVERY_NEW, clear_discovery_hash
 
@@ -73,7 +74,7 @@ async def async_setup_platform(
 ):
     """Set up MQTT switch through configuration.yaml."""
     await async_setup_reload_service(hass, DOMAIN, PLATFORMS)
-    await _async_setup_entity(hass, config, async_add_entities, discovery_info)
+    await _async_setup_entity(hass, config, async_add_entities)
 
 
 async def async_setup_entry(hass, config_entry, async_add_entities):
diff --git a/homeassistant/components/mqtt/tag.py b/homeassistant/components/mqtt/tag.py
index 94356ccf778..75f3bb50309 100644
--- a/homeassistant/components/mqtt/tag.py
+++ b/homeassistant/components/mqtt/tag.py
@@ -3,7 +3,6 @@ import logging
 
 import voluptuous as vol
 
-from homeassistant.components import mqtt
 from homeassistant.const import CONF_PLATFORM, CONF_VALUE_TEMPLATE
 import homeassistant.helpers.config_validation as cv
 from homeassistant.helpers.device_registry import EVENT_DEVICE_REGISTRY_UPDATED
@@ -21,6 +20,7 @@ from . import (
     cleanup_device_registry,
     subscription,
 )
+from .. import mqtt
 from .discovery import MQTT_DISCOVERY_NEW, MQTT_DISCOVERY_UPDATED, clear_discovery_hash
 from .util import valid_subscribe_topic
 
diff --git a/homeassistant/components/mqtt/trigger.py b/homeassistant/components/mqtt/trigger.py
index 58ec51c4b1b..1c96b3de266 100644
--- a/homeassistant/components/mqtt/trigger.py
+++ b/homeassistant/components/mqtt/trigger.py
@@ -3,11 +3,12 @@ import json
 
 import voluptuous as vol
 
-from homeassistant.components import mqtt
 from homeassistant.const import CONF_PAYLOAD, CONF_PLATFORM
 from homeassistant.core import HassJob, callback
 import homeassistant.helpers.config_validation as cv
 
+from .. import mqtt
+
 # mypy: allow-untyped-defs
 
 CONF_ENCODING = "encoding"
diff --git a/homeassistant/components/mqtt/vacuum/__init__.py b/homeassistant/components/mqtt/vacuum/__init__.py
index b954a97e8f9..f6265d1b96b 100644
--- a/homeassistant/components/mqtt/vacuum/__init__.py
+++ b/homeassistant/components/mqtt/vacuum/__init__.py
@@ -3,16 +3,12 @@ import logging
 
 import voluptuous as vol
 
-from homeassistant.components.mqtt import ATTR_DISCOVERY_HASH
-from homeassistant.components.mqtt.discovery import (
-    MQTT_DISCOVERY_NEW,
-    clear_discovery_hash,
-)
 from homeassistant.components.vacuum import DOMAIN
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.reload import async_setup_reload_service
 
-from .. import DOMAIN as MQTT_DOMAIN, PLATFORMS
+from .. import ATTR_DISCOVERY_HASH, DOMAIN as MQTT_DOMAIN, PLATFORMS
+from ..discovery import MQTT_DISCOVERY_NEW, clear_discovery_hash
 from .schema import CONF_SCHEMA, LEGACY, MQTT_VACUUM_SCHEMA, STATE
 from .schema_legacy import PLATFORM_SCHEMA_LEGACY, async_setup_entity_legacy
 from .schema_state import PLATFORM_SCHEMA_STATE, async_setup_entity_state
@@ -34,7 +30,7 @@ PLATFORM_SCHEMA = vol.All(
 async def async_setup_platform(hass, config, async_add_entities, discovery_info=None):
     """Set up MQTT vacuum through configuration.yaml."""
     await async_setup_reload_service(hass, MQTT_DOMAIN, PLATFORMS)
-    await _async_setup_entity(config, async_add_entities, discovery_info)
+    await _async_setup_entity(config, async_add_entities)
 
 
 async def async_setup_entry(hass, config_entry, async_add_entities):
@@ -58,7 +54,7 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
 
 
 async def _async_setup_entity(
-    config, async_add_entities, config_entry, discovery_data=None
+    config, async_add_entities, config_entry=None, discovery_data=None
 ):
     """Set up the MQTT vacuum."""
     setup_entity = {LEGACY: async_setup_entity_legacy, STATE: async_setup_entity_state}
diff --git a/homeassistant/components/mqtt/vacuum/schema_legacy.py b/homeassistant/components/mqtt/vacuum/schema_legacy.py
index 907e2e4a08f..65acc9afc71 100644
--- a/homeassistant/components/mqtt/vacuum/schema_legacy.py
+++ b/homeassistant/components/mqtt/vacuum/schema_legacy.py
@@ -4,14 +4,6 @@ import logging
 
 import voluptuous as vol
 
-from homeassistant.components import mqtt
-from homeassistant.components.mqtt import (
-    MqttAttributes,
-    MqttAvailability,
-    MqttDiscoveryUpdate,
-    MqttEntityDeviceInfo,
-    subscription,
-)
 from homeassistant.components.vacuum import (
     SUPPORT_BATTERY,
     SUPPORT_CLEAN_SPOT,
@@ -36,6 +28,14 @@ from homeassistant.core import callback
 import homeassistant.helpers.config_validation as cv
 from homeassistant.helpers.icon import icon_for_battery_level
 
+from .. import (
+    MqttAttributes,
+    MqttAvailability,
+    MqttDiscoveryUpdate,
+    MqttEntityDeviceInfo,
+    subscription,
+)
+from ... import mqtt
 from ..debug_info import log_messages
 from .schema import MQTT_VACUUM_SCHEMA, services_to_strings, strings_to_services
 
diff --git a/homeassistant/components/mqtt/vacuum/schema_state.py b/homeassistant/components/mqtt/vacuum/schema_state.py
index 9f75f38f1bc..5a8666e5a2e 100644
--- a/homeassistant/components/mqtt/vacuum/schema_state.py
+++ b/homeassistant/components/mqtt/vacuum/schema_state.py
@@ -4,18 +4,6 @@ import logging
 
 import voluptuous as vol
 
-from homeassistant.components import mqtt
-from homeassistant.components.mqtt import (
-    CONF_COMMAND_TOPIC,
-    CONF_QOS,
-    CONF_RETAIN,
-    CONF_STATE_TOPIC,
-    MqttAttributes,
-    MqttAvailability,
-    MqttDiscoveryUpdate,
-    MqttEntityDeviceInfo,
-    subscription,
-)
 from homeassistant.components.vacuum import (
     STATE_CLEANING,
     STATE_DOCKED,
@@ -44,6 +32,18 @@ from homeassistant.const import (
 from homeassistant.core import callback
 import homeassistant.helpers.config_validation as cv
 
+from .. import (
+    CONF_COMMAND_TOPIC,
+    CONF_QOS,
+    CONF_RETAIN,
+    CONF_STATE_TOPIC,
+    MqttAttributes,
+    MqttAvailability,
+    MqttDiscoveryUpdate,
+    MqttEntityDeviceInfo,
+    subscription,
+)
+from ... import mqtt
 from ..debug_info import log_messages
 from .schema import MQTT_VACUUM_SCHEMA, services_to_strings, strings_to_services