Fix the last failing tests

homeassistant/components/ps4/__init__.py

@@ -127,7 +127,7 @@ async def async_migrate_entry(hass, entry):
                     DOMAIN,
                     unique_id,
                     suggested_object_id=new_id,
-                    config_entry_id=e_entry.config_entry_id,
+                    config_entry=entry,
                     device_id=e_entry.device_id,
                 )
                 entry.version = 3

tests/components/ps4/test_init.py

@@ -142,7 +142,7 @@ async def test_config_flow_entry_migrate(hass):
         "media_player",
         "ps4",
         MOCK_UNIQUE_ID,
-        config_entry_id=MOCK_ENTRY_ID,
+        config_entry=mock_entry,
         device_id=MOCK_DEVICE_ID,
     )
     assert len(mock_e_registry.entities) == 1

tests/components/ps4/test_media_player.py

@@ -335,7 +335,7 @@ async def test_device_info_is_assummed(hass):
     mock_unique_id = ps4.format_unique_id(MOCK_CREDS, MOCK_HOST_ID)
     mock_e_registry = mock_registry(hass)
     mock_e_registry.async_get_or_create(
-        "media_player", DOMAIN, mock_unique_id, config_entry_id=MOCK_ENTRY_ID
+        "media_player", DOMAIN, mock_unique_id, config_entry=MOCK_CONFIG
     )
     mock_entity_id = mock_e_registry.async_get_entity_id(
         "media_player", DOMAIN, mock_unique_id

tests/helpers/test_entity_registry.py

@@ -8,7 +8,7 @@
 from homeassistant.core import valid_entity_id, callback
 from homeassistant.helpers import entity_registry
 
-from tests.common import mock_registry, flush_store
+from tests.common import MockConfigEntry, mock_registry, flush_store
 
 
 YAML__OPEN_PATH = "homeassistant.util.yaml.loader.open"
@@ -88,9 +88,11 @@ def test_create_triggers_save(hass, registry):
 
 async def test_loading_saving_data(hass, registry):
     """Test that we load/save data correctly."""
+    mock_config = MockConfigEntry(domain="light")
+
     orig_entry1 = registry.async_get_or_create("light", "hue", "1234")
     orig_entry2 = registry.async_get_or_create(
-        "light", "hue", "5678", config_entry_id="mock-id"
+        "light", "hue", "5678", config_entry=mock_config
     )
 
     assert len(registry.entities) == 2
@@ -104,7 +106,7 @@ async def test_loading_saving_data(hass, registry):
     assert list(registry.entities) == list(registry2.entities)
     new_entry1 = registry.async_get_or_create("light", "hue", "1234")
     new_entry2 = registry.async_get_or_create(
-        "light", "hue", "5678", config_entry_id="mock-id"
+        "light", "hue", "5678", config_entry=mock_config
     )
 
     assert orig_entry1 == new_entry1
@@ -198,11 +200,14 @@ def test_async_get_entity_id(registry):
 
 async def test_updating_config_entry_id(hass, registry, update_events):
     """Test that we update config entry id in registry."""
+    mock_config_1 = MockConfigEntry(domain="light", entry_id="mock-id-1")
     entry = registry.async_get_or_create(
-        "light", "hue", "5678", config_entry_id="mock-id-1"
+        "light", "hue", "5678", config_entry=mock_config_1
     )
+
+    mock_config_2 = MockConfigEntry(domain="light", entry_id="mock-id-2")
     entry2 = registry.async_get_or_create(
-        "light", "hue", "5678", config_entry_id="mock-id-2"
+        "light", "hue", "5678", config_entry=mock_config_2
     )
     assert entry.entity_id == entry2.entity_id
     assert entry2.config_entry_id == "mock-id-2"
@@ -218,8 +223,10 @@ async def test_updating_config_entry_id(hass, registry, update_events):
 
 async def test_removing_config_entry_id(hass, registry, update_events):
     """Test that we update config entry id in registry."""
+    mock_config = MockConfigEntry(domain="light", entry_id="mock-id-1")
+
     entry = registry.async_get_or_create(
-        "light", "hue", "5678", config_entry_id="mock-id-1"
+        "light", "hue", "5678", config_entry=mock_config
     )
     assert entry.config_entry_id == "mock-id-1"
     registry.async_clear_config_entry("mock-id-1")
@@ -237,6 +244,8 @@ async def test_removing_config_entry_id(hass, registry, update_events):
 
 async def test_migration(hass):
     """Test migration from old data to new."""
+    mock_config = MockConfigEntry(domain="test-platform", entry_id="test-config-id")
+
     old_conf = {
         "light.kitchen": {
             "config_entry_id": "test-config-id",
@@ -256,7 +265,7 @@ async def test_migration(hass):
         domain="light",
         platform="test-platform",
         unique_id="test-unique",
-        config_entry_id="test-config-id",
+        config_entry=mock_config,
     )
     assert entry.name == "Test Name"
     assert entry.disabled_by == "hass"
@@ -326,8 +335,10 @@ async def test_loading_race_condition(hass):
 
 async def test_update_entity_unique_id(registry):
     """Test entity's unique_id is updated."""
+    mock_config = MockConfigEntry(domain="light", entry_id="mock-id-1")
+
     entry = registry.async_get_or_create(
-        "light", "hue", "5678", config_entry_id="mock-id-1"
+        "light", "hue", "5678", config_entry=mock_config
     )
     new_unique_id = "1234"
     with patch.object(registry, "async_schedule_save") as mock_schedule_save:
@@ -341,11 +352,12 @@ async def test_update_entity_unique_id(registry):
 
 async def test_update_entity_unique_id_conflict(registry):
     """Test migration raises when unique_id already in use."""
+    mock_config = MockConfigEntry(domain="light", entry_id="mock-id-1")
     entry = registry.async_get_or_create(
-        "light", "hue", "5678", config_entry_id="mock-id-1"
+        "light", "hue", "5678", config_entry=mock_config
     )
     entry2 = registry.async_get_or_create(
-        "light", "hue", "1234", config_entry_id="mock-id-1"
+        "light", "hue", "1234", config_entry=mock_config
     )
     with patch.object(
         registry, "async_schedule_save"