From 56f2f10a17c42ba750c4cdc81ac34c0a17bc1eeb Mon Sep 17 00:00:00 2001
From: Erik Montnemery <erik@montnemery.com>
Date: Fri, 26 Apr 2024 13:49:43 +0200
Subject: [PATCH] Fix flapping trafikverket tests (#116238)

* Fix flapping trafikverket tests

* Fix copy-paste mistake
---
 .../components/trafikverket_train/conftest.py | 38 +++++++++++--------
 1 file changed, 22 insertions(+), 16 deletions(-)

diff --git a/tests/components/trafikverket_train/conftest.py b/tests/components/trafikverket_train/conftest.py
index 880701e8bdc..7221d96bae2 100644
--- a/tests/components/trafikverket_train/conftest.py
+++ b/tests/components/trafikverket_train/conftest.py
@@ -25,6 +25,25 @@ async def load_integration_from_entry(
     get_train_stop: TrainStop,
 ) -> MockConfigEntry:
     """Set up the Trafikverket Train integration in Home Assistant."""
+
+    async def setup_config_entry_with_mocked_data(config_entry_id: str) -> None:
+        """Set up a config entry with mocked trafikverket data."""
+        with (
+            patch(
+                "homeassistant.components.trafikverket_train.coordinator.TrafikverketTrain.async_get_next_train_stops",
+                return_value=get_trains,
+            ),
+            patch(
+                "homeassistant.components.trafikverket_train.coordinator.TrafikverketTrain.async_get_train_stop",
+                return_value=get_train_stop,
+            ),
+            patch(
+                "homeassistant.components.trafikverket_train.TrafikverketTrain.async_get_train_station",
+            ),
+        ):
+            await hass.config_entries.async_setup(config_entry_id)
+            await hass.async_block_till_done()
+
     config_entry = MockConfigEntry(
         domain=DOMAIN,
         source=SOURCE_USER,
@@ -34,6 +53,8 @@ async def load_integration_from_entry(
         unique_id="stockholmc-uppsalac--['mon', 'tue', 'wed', 'thu', 'fri', 'sat', 'sun']",
     )
     config_entry.add_to_hass(hass)
+    await setup_config_entry_with_mocked_data(config_entry.entry_id)
+
     config_entry2 = MockConfigEntry(
         domain=DOMAIN,
         source=SOURCE_USER,
@@ -42,22 +63,7 @@ async def load_integration_from_entry(
         unique_id="stockholmc-uppsalac-1100-['mon', 'tue', 'wed', 'thu', 'fri', 'sat', 'sun']",
     )
     config_entry2.add_to_hass(hass)
-
-    with (
-        patch(
-            "homeassistant.components.trafikverket_train.coordinator.TrafikverketTrain.async_get_next_train_stops",
-            return_value=get_trains,
-        ),
-        patch(
-            "homeassistant.components.trafikverket_train.coordinator.TrafikverketTrain.async_get_train_stop",
-            return_value=get_train_stop,
-        ),
-        patch(
-            "homeassistant.components.trafikverket_train.TrafikverketTrain.async_get_train_station",
-        ),
-    ):
-        await hass.config_entries.async_setup(config_entry.entry_id)
-        await hass.async_block_till_done()
+    await setup_config_entry_with_mocked_data(config_entry2.entry_id)
 
     return config_entry