diff --git a/tests/components/axis/test_binary_sensor.py b/tests/components/axis/test_binary_sensor.py
index e2d820a959e..0c6f4535cd0 100644
--- a/tests/components/axis/test_binary_sensor.py
+++ b/tests/components/axis/test_binary_sensor.py
@@ -5,6 +5,7 @@ from homeassistant.components.binary_sensor import (
     DEVICE_CLASS_MOTION,
     DOMAIN as BINARY_SENSOR_DOMAIN,
 )
+from homeassistant.const import STATE_OFF, STATE_ON
 from homeassistant.setup import async_setup_component
 
 from .test_device import NAME, setup_axis_integration
@@ -50,7 +51,8 @@ async def test_no_binary_sensors(hass):
 
 async def test_binary_sensors(hass):
     """Test that sensors are loaded properly."""
-    device = await setup_axis_integration(hass)
+    config_entry = await setup_axis_integration(hass)
+    device = hass.data[AXIS_DOMAIN][config_entry.unique_id]
 
     for event in EVENTS:
         device.api.event.process_event(event)
@@ -58,12 +60,12 @@ async def test_binary_sensors(hass):
 
     assert len(hass.states.async_entity_ids(BINARY_SENSOR_DOMAIN)) == 2
 
-    pir = hass.states.get(f"binary_sensor.{NAME}_pir_0")
-    assert pir.state == "off"
+    pir = hass.states.get(f"{BINARY_SENSOR_DOMAIN}.{NAME}_pir_0")
+    assert pir.state == STATE_OFF
     assert pir.name == f"{NAME} PIR 0"
     assert pir.attributes["device_class"] == DEVICE_CLASS_MOTION
 
-    vmd4 = hass.states.get(f"binary_sensor.{NAME}_vmd4_profile_1")
-    assert vmd4.state == "on"
+    vmd4 = hass.states.get(f"{BINARY_SENSOR_DOMAIN}.{NAME}_vmd4_profile_1")
+    assert vmd4.state == STATE_ON
     assert vmd4.name == f"{NAME} VMD4 Profile 1"
     assert vmd4.attributes["device_class"] == DEVICE_CLASS_MOTION
diff --git a/tests/components/axis/test_camera.py b/tests/components/axis/test_camera.py
index af276fe6fe5..01c5a677e38 100644
--- a/tests/components/axis/test_camera.py
+++ b/tests/components/axis/test_camera.py
@@ -6,6 +6,7 @@ from homeassistant.components.axis.const import (
     DOMAIN as AXIS_DOMAIN,
 )
 from homeassistant.components.camera import DOMAIN as CAMERA_DOMAIN
+from homeassistant.const import STATE_IDLE
 from homeassistant.setup import async_setup_component
 
 from .test_device import ENTRY_OPTIONS, NAME, setup_axis_integration
@@ -17,7 +18,7 @@ async def test_platform_manually_configured(hass):
     """Test that nothing happens when platform is manually configured."""
     assert (
         await async_setup_component(
-            hass, CAMERA_DOMAIN, {"camera": {"platform": AXIS_DOMAIN}}
+            hass, CAMERA_DOMAIN, {CAMERA_DOMAIN: {"platform": AXIS_DOMAIN}}
         )
         is True
     )
@@ -31,11 +32,13 @@ async def test_camera(hass):
 
     assert len(hass.states.async_entity_ids(CAMERA_DOMAIN)) == 1
 
-    cam = hass.states.get(f"camera.{NAME}")
-    assert cam.state == "idle"
+    entity_id = f"{CAMERA_DOMAIN}.{NAME}"
+
+    cam = hass.states.get(entity_id)
+    assert cam.state == STATE_IDLE
     assert cam.name == NAME
 
-    camera_entity = camera._get_camera_from_entity_id(hass, f"camera.{NAME}")
+    camera_entity = camera._get_camera_from_entity_id(hass, entity_id)
     assert camera_entity.image_source == "http://1.2.3.4:80/axis-cgi/jpg/image.cgi"
     assert camera_entity.mjpeg_source == "http://1.2.3.4:80/axis-cgi/mjpg/video.cgi"
     assert (
@@ -51,11 +54,13 @@ async def test_camera_with_stream_profile(hass):
 
     assert len(hass.states.async_entity_ids(CAMERA_DOMAIN)) == 1
 
-    cam = hass.states.get(f"camera.{NAME}")
-    assert cam.state == "idle"
+    entity_id = f"{CAMERA_DOMAIN}.{NAME}"
+
+    cam = hass.states.get(entity_id)
+    assert cam.state == STATE_IDLE
     assert cam.name == NAME
 
-    camera_entity = camera._get_camera_from_entity_id(hass, f"camera.{NAME}")
+    camera_entity = camera._get_camera_from_entity_id(hass, entity_id)
     assert camera_entity.image_source == "http://1.2.3.4:80/axis-cgi/jpg/image.cgi"
     assert (
         camera_entity.mjpeg_source
diff --git a/tests/components/axis/test_config_flow.py b/tests/components/axis/test_config_flow.py
index 0def1d84251..24f888ea6ef 100644
--- a/tests/components/axis/test_config_flow.py
+++ b/tests/components/axis/test_config_flow.py
@@ -8,6 +8,7 @@ from homeassistant.components.axis.const import (
     DEFAULT_STREAM_PROFILE,
     DOMAIN as AXIS_DOMAIN,
 )
+from homeassistant.config_entries import SOURCE_USER, SOURCE_ZEROCONF
 from homeassistant.const import (
     CONF_HOST,
     CONF_MAC,
@@ -16,6 +17,11 @@ from homeassistant.const import (
     CONF_PORT,
     CONF_USERNAME,
 )
+from homeassistant.data_entry_flow import (
+    RESULT_TYPE_ABORT,
+    RESULT_TYPE_CREATE_ENTRY,
+    RESULT_TYPE_FORM,
+)
 
 from .test_device import MAC, MODEL, NAME, setup_axis_integration, vapix_request
 
@@ -26,11 +32,11 @@ from tests.common import MockConfigEntry
 async def test_flow_manual_configuration(hass):
     """Test that config flow works."""
     result = await hass.config_entries.flow.async_init(
-        AXIS_DOMAIN, context={"source": "user"}
+        AXIS_DOMAIN, context={"source": SOURCE_USER}
     )
 
-    assert result["type"] == "form"
-    assert result["step_id"] == "user"
+    assert result["type"] == RESULT_TYPE_FORM
+    assert result["step_id"] == SOURCE_USER
 
     with patch("axis.vapix.Vapix.request", new=vapix_request):
         result = await hass.config_entries.flow.async_configure(
@@ -43,7 +49,7 @@ async def test_flow_manual_configuration(hass):
             },
         )
 
-    assert result["type"] == "create_entry"
+    assert result["type"] == RESULT_TYPE_CREATE_ENTRY
     assert result["title"] == f"M1065-LW - {MAC}"
     assert result["data"] == {
         CONF_HOST: "1.2.3.4",
@@ -58,14 +64,15 @@ async def test_flow_manual_configuration(hass):
 
 async def test_manual_configuration_update_configuration(hass):
     """Test that config flow fails on already configured device."""
-    device = await setup_axis_integration(hass)
+    config_entry = await setup_axis_integration(hass)
+    device = hass.data[AXIS_DOMAIN][config_entry.unique_id]
 
     result = await hass.config_entries.flow.async_init(
-        AXIS_DOMAIN, context={"source": "user"}
+        AXIS_DOMAIN, context={"source": SOURCE_USER}
     )
 
-    assert result["type"] == "form"
-    assert result["step_id"] == "user"
+    assert result["type"] == RESULT_TYPE_FORM
+    assert result["step_id"] == SOURCE_USER
 
     with patch(
         "homeassistant.components.axis.async_setup_entry",
@@ -82,7 +89,7 @@ async def test_manual_configuration_update_configuration(hass):
         )
         await hass.async_block_till_done()
 
-    assert result["type"] == "abort"
+    assert result["type"] == RESULT_TYPE_ABORT
     assert result["reason"] == "already_configured"
     assert device.host == "2.3.4.5"
     assert len(mock_setup_entry.mock_calls) == 1
@@ -93,11 +100,11 @@ async def test_flow_fails_already_configured(hass):
     await setup_axis_integration(hass)
 
     result = await hass.config_entries.flow.async_init(
-        AXIS_DOMAIN, context={"source": "user"}
+        AXIS_DOMAIN, context={"source": SOURCE_USER}
     )
 
-    assert result["type"] == "form"
-    assert result["step_id"] == "user"
+    assert result["type"] == RESULT_TYPE_FORM
+    assert result["step_id"] == SOURCE_USER
 
     with patch("axis.vapix.Vapix.request", new=vapix_request):
         result = await hass.config_entries.flow.async_configure(
@@ -110,18 +117,18 @@ async def test_flow_fails_already_configured(hass):
             },
         )
 
-    assert result["type"] == "abort"
+    assert result["type"] == RESULT_TYPE_ABORT
     assert result["reason"] == "already_configured"
 
 
 async def test_flow_fails_faulty_credentials(hass):
     """Test that config flow fails on faulty credentials."""
     result = await hass.config_entries.flow.async_init(
-        AXIS_DOMAIN, context={"source": "user"}
+        AXIS_DOMAIN, context={"source": SOURCE_USER}
     )
 
-    assert result["type"] == "form"
-    assert result["step_id"] == "user"
+    assert result["type"] == RESULT_TYPE_FORM
+    assert result["step_id"] == SOURCE_USER
 
     with patch(
         "homeassistant.components.axis.config_flow.get_device",
@@ -143,11 +150,11 @@ async def test_flow_fails_faulty_credentials(hass):
 async def test_flow_fails_cannot_connect(hass):
     """Test that config flow fails on cannot connect."""
     result = await hass.config_entries.flow.async_init(
-        AXIS_DOMAIN, context={"source": "user"}
+        AXIS_DOMAIN, context={"source": SOURCE_USER}
     )
 
-    assert result["type"] == "form"
-    assert result["step_id"] == "user"
+    assert result["type"] == RESULT_TYPE_FORM
+    assert result["step_id"] == SOURCE_USER
 
     with patch(
         "homeassistant.components.axis.config_flow.get_device",
@@ -180,11 +187,11 @@ async def test_flow_create_entry_multiple_existing_entries_of_same_model(hass):
     entry2.add_to_hass(hass)
 
     result = await hass.config_entries.flow.async_init(
-        AXIS_DOMAIN, context={"source": "user"}
+        AXIS_DOMAIN, context={"source": SOURCE_USER}
     )
 
-    assert result["type"] == "form"
-    assert result["step_id"] == "user"
+    assert result["type"] == RESULT_TYPE_FORM
+    assert result["step_id"] == SOURCE_USER
 
     with patch("axis.vapix.Vapix.request", new=vapix_request):
         result = await hass.config_entries.flow.async_configure(
@@ -197,7 +204,7 @@ async def test_flow_create_entry_multiple_existing_entries_of_same_model(hass):
             },
         )
 
-    assert result["type"] == "create_entry"
+    assert result["type"] == RESULT_TYPE_CREATE_ENTRY
     assert result["title"] == f"M1065-LW - {MAC}"
     assert result["data"] == {
         CONF_HOST: "1.2.3.4",
@@ -222,11 +229,11 @@ async def test_zeroconf_flow(hass):
             "hostname": "name",
             "properties": {"macaddress": MAC},
         },
-        context={"source": "zeroconf"},
+        context={"source": SOURCE_ZEROCONF},
     )
 
-    assert result["type"] == "form"
-    assert result["step_id"] == "user"
+    assert result["type"] == RESULT_TYPE_FORM
+    assert result["step_id"] == SOURCE_USER
 
     with patch("axis.vapix.Vapix.request", new=vapix_request):
         result = await hass.config_entries.flow.async_configure(
@@ -239,7 +246,7 @@ async def test_zeroconf_flow(hass):
             },
         )
 
-    assert result["type"] == "create_entry"
+    assert result["type"] == RESULT_TYPE_CREATE_ENTRY
     assert result["title"] == f"M1065-LW - {MAC}"
     assert result["data"] == {
         CONF_HOST: "1.2.3.4",
@@ -256,7 +263,8 @@ async def test_zeroconf_flow(hass):
 
 async def test_zeroconf_flow_already_configured(hass):
     """Test that zeroconf doesn't setup already configured devices."""
-    device = await setup_axis_integration(hass)
+    config_entry = await setup_axis_integration(hass)
+    device = hass.data[AXIS_DOMAIN][config_entry.unique_id]
     assert device.host == "1.2.3.4"
 
     result = await hass.config_entries.flow.async_init(
@@ -267,17 +275,18 @@ async def test_zeroconf_flow_already_configured(hass):
             "hostname": "name",
             "properties": {"macaddress": MAC},
         },
-        context={"source": "zeroconf"},
+        context={"source": SOURCE_ZEROCONF},
     )
 
-    assert result["type"] == "abort"
+    assert result["type"] == RESULT_TYPE_ABORT
     assert result["reason"] == "already_configured"
     assert device.host == "1.2.3.4"
 
 
 async def test_zeroconf_flow_updated_configuration(hass):
     """Test that zeroconf update configuration with new parameters."""
-    device = await setup_axis_integration(hass)
+    config_entry = await setup_axis_integration(hass)
+    device = hass.data[AXIS_DOMAIN][config_entry.unique_id]
     assert device.host == "1.2.3.4"
     assert device.config_entry.data == {
         CONF_HOST: "1.2.3.4",
@@ -301,11 +310,11 @@ async def test_zeroconf_flow_updated_configuration(hass):
                 "hostname": "name",
                 "properties": {"macaddress": MAC},
             },
-            context={"source": "zeroconf"},
+            context={"source": SOURCE_ZEROCONF},
         )
         await hass.async_block_till_done()
 
-    assert result["type"] == "abort"
+    assert result["type"] == RESULT_TYPE_ABORT
     assert result["reason"] == "already_configured"
     assert device.config_entry.data == {
         CONF_HOST: "2.3.4.5",
@@ -324,10 +333,10 @@ async def test_zeroconf_flow_ignore_non_axis_device(hass):
     result = await hass.config_entries.flow.async_init(
         AXIS_DOMAIN,
         data={CONF_HOST: "169.254.3.4", "properties": {"macaddress": "01234567890"}},
-        context={"source": "zeroconf"},
+        context={"source": SOURCE_ZEROCONF},
     )
 
-    assert result["type"] == "abort"
+    assert result["type"] == RESULT_TYPE_ABORT
     assert result["reason"] == "not_axis_device"
 
 
@@ -336,16 +345,17 @@ async def test_zeroconf_flow_ignore_link_local_address(hass):
     result = await hass.config_entries.flow.async_init(
         AXIS_DOMAIN,
         data={CONF_HOST: "169.254.3.4", "properties": {"macaddress": MAC}},
-        context={"source": "zeroconf"},
+        context={"source": SOURCE_ZEROCONF},
     )
 
-    assert result["type"] == "abort"
+    assert result["type"] == RESULT_TYPE_ABORT
     assert result["reason"] == "link_local_address"
 
 
 async def test_option_flow(hass):
     """Test config flow options."""
-    device = await setup_axis_integration(hass)
+    config_entry = await setup_axis_integration(hass)
+    device = hass.data[AXIS_DOMAIN][config_entry.unique_id]
     assert device.option_stream_profile == DEFAULT_STREAM_PROFILE
 
     result = await hass.config_entries.options.async_init(device.config_entry.entry_id)
diff --git a/tests/components/axis/test_device.py b/tests/components/axis/test_device.py
index 627a17ecf60..b7612a01c1d 100644
--- a/tests/components/axis/test_device.py
+++ b/tests/components/axis/test_device.py
@@ -29,6 +29,7 @@ from homeassistant.components.axis.const import (
     DOMAIN as AXIS_DOMAIN,
 )
 from homeassistant.components.binary_sensor import DOMAIN as BINARY_SENSOR_DOMAIN
+from homeassistant.config_entries import SOURCE_ZEROCONF
 from homeassistant.const import (
     CONF_HOST,
     CONF_MAC,
@@ -36,6 +37,7 @@ from homeassistant.const import (
     CONF_PASSWORD,
     CONF_PORT,
     CONF_USERNAME,
+    STATE_ON,
 )
 
 from tests.async_mock import AsyncMock, Mock, patch
@@ -229,7 +231,6 @@ async def setup_axis_integration(hass, config=ENTRY_CONFIG, options=ENTRY_OPTION
         data=deepcopy(config),
         connection_class=config_entries.CONN_CLASS_LOCAL_PUSH,
         options=deepcopy(options),
-        entry_id="1",
         version=2,
     )
     config_entry.add_to_hass(hass)
@@ -241,7 +242,7 @@ async def setup_axis_integration(hass, config=ENTRY_CONFIG, options=ENTRY_OPTION
         await hass.config_entries.async_setup(config_entry.entry_id)
         await hass.async_block_till_done()
 
-    return hass.data[AXIS_DOMAIN].get(config_entry.unique_id)
+    return config_entry
 
 
 async def test_device_setup(hass):
@@ -250,7 +251,8 @@ async def test_device_setup(hass):
         "homeassistant.config_entries.ConfigEntries.async_forward_entry_setup",
         return_value=True,
     ) as forward_entry_setup:
-        device = await setup_axis_integration(hass)
+        config_entry = await setup_axis_integration(hass)
+        device = hass.data[AXIS_DOMAIN][config_entry.unique_id]
 
     assert device.api.vapix.firmware_version == "9.10.1"
     assert device.api.vapix.product_number == "M1065-LW"
@@ -277,7 +279,8 @@ async def test_device_info(hass):
     api_discovery["data"]["apiList"].append(API_DISCOVERY_BASIC_DEVICE_INFO)
 
     with patch.dict(API_DISCOVERY_RESPONSE, api_discovery):
-        device = await setup_axis_integration(hass)
+        config_entry = await setup_axis_integration(hass)
+        device = hass.data[AXIS_DOMAIN][config_entry.unique_id]
 
     assert device.api.vapix.firmware_version == "9.80.1"
     assert device.api.vapix.product_number == "M1065-LW"
@@ -303,14 +306,15 @@ async def test_device_support_mqtt(hass, mqtt_mock):
     await hass.async_block_till_done()
     assert len(hass.states.async_entity_ids(BINARY_SENSOR_DOMAIN)) == 1
 
-    pir = hass.states.get(f"binary_sensor.{NAME}_pir_0")
-    assert pir.state == "on"
+    pir = hass.states.get(f"{BINARY_SENSOR_DOMAIN}.{NAME}_pir_0")
+    assert pir.state == STATE_ON
     assert pir.name == f"{NAME} PIR 0"
 
 
 async def test_update_address(hass):
     """Test update address works."""
-    device = await setup_axis_integration(hass)
+    config_entry = await setup_axis_integration(hass)
+    device = hass.data[AXIS_DOMAIN][config_entry.unique_id]
     assert device.api.config.host == "1.2.3.4"
 
     with patch("axis.vapix.Vapix.request", new=vapix_request), patch(
@@ -325,7 +329,7 @@ async def test_update_address(hass):
                 "hostname": "name",
                 "properties": {"macaddress": MAC},
             },
-            context={"source": "zeroconf"},
+            context={"source": SOURCE_ZEROCONF},
         )
         await hass.async_block_till_done()
 
@@ -335,14 +339,16 @@ async def test_update_address(hass):
 
 async def test_device_unavailable(hass):
     """Successful setup."""
-    device = await setup_axis_integration(hass)
+    config_entry = await setup_axis_integration(hass)
+    device = hass.data[AXIS_DOMAIN][config_entry.unique_id]
     device.async_connection_status_callback(status=False)
     assert not device.available
 
 
 async def test_device_reset(hass):
     """Successfully reset device."""
-    device = await setup_axis_integration(hass)
+    config_entry = await setup_axis_integration(hass)
+    device = hass.data[AXIS_DOMAIN][config_entry.unique_id]
     result = await device.async_reset()
     assert result is True
 
diff --git a/tests/components/axis/test_init.py b/tests/components/axis/test_init.py
index 3feee94267a..cf5253d4675 100644
--- a/tests/components/axis/test_init.py
+++ b/tests/components/axis/test_init.py
@@ -51,7 +51,8 @@ async def test_setup_entry_fails(hass):
 
 async def test_unload_entry(hass):
     """Test successful unload of entry."""
-    device = await setup_axis_integration(hass)
+    config_entry = await setup_axis_integration(hass)
+    device = hass.data[AXIS_DOMAIN][config_entry.unique_id]
     assert hass.data[AXIS_DOMAIN]
 
     assert await hass.config_entries.async_unload(device.config_entry.entry_id)
diff --git a/tests/components/axis/test_light.py b/tests/components/axis/test_light.py
index 987ef1cd48b..06612daa313 100644
--- a/tests/components/axis/test_light.py
+++ b/tests/components/axis/test_light.py
@@ -4,6 +4,13 @@ from copy import deepcopy
 
 from homeassistant.components.axis.const import DOMAIN as AXIS_DOMAIN
 from homeassistant.components.light import ATTR_BRIGHTNESS, DOMAIN as LIGHT_DOMAIN
+from homeassistant.const import (
+    ATTR_ENTITY_ID,
+    SERVICE_TURN_OFF,
+    SERVICE_TURN_ON,
+    STATE_OFF,
+    STATE_ON,
+)
 from homeassistant.setup import async_setup_component
 
 from .test_device import API_DISCOVERY_RESPONSE, NAME, setup_axis_integration
@@ -57,7 +64,8 @@ async def test_lights(hass):
     api_discovery["data"]["apiList"].append(API_DISCOVERY_LIGHT_CONTROL)
 
     with patch.dict(API_DISCOVERY_RESPONSE, api_discovery):
-        device = await setup_axis_integration(hass)
+        config_entry = await setup_axis_integration(hass)
+        device = hass.data[AXIS_DOMAIN][config_entry.unique_id]
 
     # Add light
     with patch(
@@ -72,8 +80,10 @@ async def test_lights(hass):
 
     assert len(hass.states.async_entity_ids(LIGHT_DOMAIN)) == 1
 
-    light_0 = hass.states.get(f"light.{NAME}_ir_light_0")
-    assert light_0.state == "on"
+    entity_id = f"{LIGHT_DOMAIN}.{NAME}_ir_light_0"
+
+    light_0 = hass.states.get(entity_id)
+    assert light_0.state == STATE_ON
     assert light_0.name == f"{NAME} IR Light 0"
 
     # Turn on, set brightness, light already on
@@ -87,8 +97,8 @@ async def test_lights(hass):
     ):
         await hass.services.async_call(
             LIGHT_DOMAIN,
-            "turn_on",
-            {"entity_id": f"light.{NAME}_ir_light_0", ATTR_BRIGHTNESS: 50},
+            SERVICE_TURN_ON,
+            {ATTR_ENTITY_ID: entity_id, ATTR_BRIGHTNESS: 50},
             blocking=True,
         )
         mock_activate.assert_not_awaited()
@@ -103,8 +113,8 @@ async def test_lights(hass):
     ):
         await hass.services.async_call(
             LIGHT_DOMAIN,
-            "turn_off",
-            {"entity_id": f"light.{NAME}_ir_light_0"},
+            SERVICE_TURN_OFF,
+            {ATTR_ENTITY_ID: entity_id},
             blocking=True,
         )
         mock_deactivate.assert_called_once()
@@ -113,8 +123,8 @@ async def test_lights(hass):
     device.api.event.process_event(EVENT_OFF)
     await hass.async_block_till_done()
 
-    light_0 = hass.states.get(f"light.{NAME}_ir_light_0")
-    assert light_0.state == "off"
+    light_0 = hass.states.get(entity_id)
+    assert light_0.state == STATE_OFF
 
     # Turn on, set brightness
     with patch(
@@ -127,8 +137,8 @@ async def test_lights(hass):
     ):
         await hass.services.async_call(
             LIGHT_DOMAIN,
-            "turn_on",
-            {"entity_id": f"light.{NAME}_ir_light_0"},
+            SERVICE_TURN_ON,
+            {ATTR_ENTITY_ID: entity_id},
             blocking=True,
         )
         mock_activate.assert_called_once()
@@ -143,8 +153,8 @@ async def test_lights(hass):
     ):
         await hass.services.async_call(
             LIGHT_DOMAIN,
-            "turn_off",
-            {"entity_id": f"light.{NAME}_ir_light_0"},
+            SERVICE_TURN_OFF,
+            {ATTR_ENTITY_ID: entity_id},
             blocking=True,
         )
         mock_deactivate.assert_not_called()
diff --git a/tests/components/axis/test_switch.py b/tests/components/axis/test_switch.py
index 37d50e33e3e..fbcf0624fc9 100644
--- a/tests/components/axis/test_switch.py
+++ b/tests/components/axis/test_switch.py
@@ -4,6 +4,13 @@ from copy import deepcopy
 
 from homeassistant.components.axis.const import DOMAIN as AXIS_DOMAIN
 from homeassistant.components.switch import DOMAIN as SWITCH_DOMAIN
+from homeassistant.const import (
+    ATTR_ENTITY_ID,
+    SERVICE_TURN_OFF,
+    SERVICE_TURN_ON,
+    STATE_OFF,
+    STATE_ON,
+)
 from homeassistant.setup import async_setup_component
 
 from .test_device import (
@@ -53,7 +60,8 @@ async def test_no_switches(hass):
 
 async def test_switches_with_port_cgi(hass):
     """Test that switches are loaded properly using port.cgi."""
-    device = await setup_axis_integration(hass)
+    config_entry = await setup_axis_integration(hass)
+    device = hass.data[AXIS_DOMAIN][config_entry.unique_id]
 
     device.api.vapix.ports = {"0": AsyncMock(), "1": AsyncMock()}
     device.api.vapix.ports["0"].name = "Doorbell"
@@ -67,26 +75,28 @@ async def test_switches_with_port_cgi(hass):
 
     assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 2
 
-    relay_0 = hass.states.get(f"switch.{NAME}_doorbell")
-    assert relay_0.state == "off"
-    assert relay_0.name == f"{NAME} Doorbell"
-
-    relay_1 = hass.states.get(f"switch.{NAME}_relay_1")
-    assert relay_1.state == "on"
+    relay_1 = hass.states.get(f"{SWITCH_DOMAIN}.{NAME}_relay_1")
+    assert relay_1.state == STATE_ON
     assert relay_1.name == f"{NAME} Relay 1"
 
+    entity_id = f"{SWITCH_DOMAIN}.{NAME}_doorbell"
+
+    relay_0 = hass.states.get(entity_id)
+    assert relay_0.state == STATE_OFF
+    assert relay_0.name == f"{NAME} Doorbell"
+
     await hass.services.async_call(
         SWITCH_DOMAIN,
-        "turn_on",
-        {"entity_id": f"switch.{NAME}_doorbell"},
+        SERVICE_TURN_ON,
+        {ATTR_ENTITY_ID: entity_id},
         blocking=True,
     )
     device.api.vapix.ports["0"].close.assert_called_once()
 
     await hass.services.async_call(
         SWITCH_DOMAIN,
-        "turn_off",
-        {"entity_id": f"switch.{NAME}_doorbell"},
+        SERVICE_TURN_OFF,
+        {ATTR_ENTITY_ID: entity_id},
         blocking=True,
     )
     device.api.vapix.ports["0"].open.assert_called_once()
@@ -98,7 +108,8 @@ async def test_switches_with_port_management(hass):
     api_discovery["data"]["apiList"].append(API_DISCOVERY_PORT_MANAGEMENT)
 
     with patch.dict(API_DISCOVERY_RESPONSE, api_discovery):
-        device = await setup_axis_integration(hass)
+        config_entry = await setup_axis_integration(hass)
+        device = hass.data[AXIS_DOMAIN][config_entry.unique_id]
 
     device.api.vapix.ports = {"0": AsyncMock(), "1": AsyncMock()}
     device.api.vapix.ports["0"].name = "Doorbell"
@@ -112,26 +123,28 @@ async def test_switches_with_port_management(hass):
 
     assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 2
 
-    relay_0 = hass.states.get(f"switch.{NAME}_doorbell")
-    assert relay_0.state == "off"
-    assert relay_0.name == f"{NAME} Doorbell"
-
-    relay_1 = hass.states.get(f"switch.{NAME}_relay_1")
-    assert relay_1.state == "on"
+    relay_1 = hass.states.get(f"{SWITCH_DOMAIN}.{NAME}_relay_1")
+    assert relay_1.state == STATE_ON
     assert relay_1.name == f"{NAME} Relay 1"
 
+    entity_id = f"{SWITCH_DOMAIN}.{NAME}_doorbell"
+
+    relay_0 = hass.states.get(entity_id)
+    assert relay_0.state == STATE_OFF
+    assert relay_0.name == f"{NAME} Doorbell"
+
     await hass.services.async_call(
         SWITCH_DOMAIN,
-        "turn_on",
-        {"entity_id": f"switch.{NAME}_doorbell"},
+        SERVICE_TURN_ON,
+        {ATTR_ENTITY_ID: entity_id},
         blocking=True,
     )
     device.api.vapix.ports["0"].close.assert_called_once()
 
     await hass.services.async_call(
         SWITCH_DOMAIN,
-        "turn_off",
-        {"entity_id": f"switch.{NAME}_doorbell"},
+        SERVICE_TURN_OFF,
+        {ATTR_ENTITY_ID: entity_id},
         blocking=True,
     )
     device.api.vapix.ports["0"].open.assert_called_once()