Feat: support multiple Integration entries

custom_components/flexmeasures_hacs/__init__.py

@@ -31,8 +31,6 @@
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Set up FlexMeasures from a config entry."""
 
-    hass.data.setdefault(DOMAIN, {})
-
     # Reload integration when the options are updated
     entry.async_on_unload(entry.add_update_listener(options_update_listener))
 
@@ -69,9 +67,11 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
             )
 
     FRBC_data = FRBC_Config(**frbc_data_dict)
-    hass.data[DOMAIN][FRBC_CONFIG] = FRBC_data
-
-    hass.data[DOMAIN]["fm_client"] = client
+    entry_config = {
+        FRBC_CONFIG: FRBC_data,
+        "fm_client": client,
+    }
+    hass.data.setdefault(DOMAIN, {})[entry.entry_id] = entry_config
 
     hass.http.register_view(WebsocketAPIView(entry))
 
@@ -98,7 +98,7 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     await async_unload_services(hass)
 
     if unload_ok := await hass.config_entries.async_unload_platforms(entry, PLATFORMS):
-        hass.data[DOMAIN].pop(entry.entry_id, None)
+        hass.data[DOMAIN].pop(entry.entry_id)
 
     return unload_ok
 

custom_components/flexmeasures_hacs/sensor.py

@@ -21,9 +21,9 @@ async def async_setup_entry(
     hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
 ) -> None:
     """Set up sensor."""
-    hass.data[DOMAIN][SCHEDULE_STATE] = {"schedule": [], "start": None}
+    hass.data[DOMAIN][entry.entry_id][SCHEDULE_STATE] = {"schedule": [], "start": None}
 
-    async_add_entities([FlexMeasuresScheduleSensor()], True)
+    async_add_entities([FlexMeasuresScheduleSensor(entry_id=entry.entry_id)], True)
 
 
 class FlexMeasuresScheduleSensor(SensorEntity):
@@ -32,9 +32,10 @@ class FlexMeasuresScheduleSensor(SensorEntity):
     _attr_device_class = SensorDeviceClass.POWER
     _attr_native_unit_of_measurement = UnitOfPower.KILO_WATT
 
-    def __init__(self) -> None:
+    def __init__(self, entry_id) -> None:
         """Sensor to store the schedule created by FlexMeasures."""
         self._attr_unique_id = SCHEDULE_ENTITY
+        self.entry_id = entry_id
 
     @property
     def name(self) -> str:
@@ -45,15 +46,15 @@ def name(self) -> str:
     def native_value(self) -> float:
         """Average power."""
 
-        commands = self.hass.data[DOMAIN][SCHEDULE_STATE]["schedule"]
+        commands = self.hass.data[DOMAIN][self.entry_id][SCHEDULE_STATE]["schedule"]
         if len(commands) == 0:
             return 0
         return sum(command["value"] for command in commands) / len(commands)
 
     @property
     def extra_state_attributes(self) -> dict[str, Any]:
         """Return default attributes for the FlexMeasures Schedule sensor."""
-        return self.hass.data[DOMAIN][SCHEDULE_STATE]
+        return self.hass.data[DOMAIN][self.entry_id][SCHEDULE_STATE]
 
     async def async_added_to_hass(self) -> None:
         """Register callbacks."""

custom_components/flexmeasures_hacs/services.py

@@ -81,10 +81,10 @@ async def change_control_type(
     ):  # pylint: disable=possibly-unused-variable
         """Change control type S2 Protocol."""
 
-        if "cem" not in hass.data[DOMAIN]:
+        if "cem" not in hass.data[DOMAIN][entry.entry_id]:
             raise UndefinedCEMError()
 
-        cem: CEM = hass.data[DOMAIN]["cem"]
+        cem: CEM = hass.data[DOMAIN][entry.entry_id]["cem"]
 
         control_type = cast(str, call.data.get("control_type"))
 
@@ -96,13 +96,14 @@ async def change_control_type(
         await cem.activate_control_type(control_type=control_type)
 
         hass.states.async_set(
-            f"{DOMAIN}.cem", json.dumps({"control_type": str(cem.control_type)})
+            f"{DOMAIN}.{entry.entry_id}.cem",
+            json.dumps({"control_type": str(cem.control_type)}),
         )
 
     async def trigger_and_get_schedule(
         call: ServiceCall,
     ):  # pylint: disable=possibly-unused-variable
-        client: FlexMeasuresClient = hass.data[DOMAIN]["fm_client"]
+        client: FlexMeasuresClient = hass.data[DOMAIN][entry.entry_id]["fm_client"]
         resolution = pd.Timedelta(RESOLUTION)
         tzinfo = dt_util.get_time_zone(hass.config.time_zone)
         start = time_ceil(datetime.now(tz=tzinfo), resolution)
@@ -146,18 +147,18 @@ async def trigger_and_get_schedule(
             for i, value in enumerate(schedule["values"])
         ]
 
-        hass.data[DOMAIN][SCHEDULE_STATE]["schedule"] = schedule
-        hass.data[DOMAIN][SCHEDULE_STATE]["start"] = start
-        hass.data[DOMAIN][SCHEDULE_STATE]["duration"] = get_from_option_or_config(
-            "schedule_duration", entry
+        hass.data[DOMAIN][entry.entry_id][SCHEDULE_STATE]["schedule"] = schedule
+        hass.data[DOMAIN][entry.entry_id][SCHEDULE_STATE]["start"] = start
+        hass.data[DOMAIN][entry.entry_id][SCHEDULE_STATE]["duration"] = (
+            get_from_option_or_config("schedule_duration", entry)
         )
 
         async_dispatcher_send(hass, SIGNAL_UPDATE_SCHEDULE)
 
     async def post_measurements(
         call: ServiceCall,
     ):  # pylint: disable=possibly-unused-variable
-        client: FlexMeasuresClient = hass.data[DOMAIN]["fm_client"]
+        client: FlexMeasuresClient = hass.data[DOMAIN][entry.entry_id]["fm_client"]
 
         await client.post_measurements(
             sensor_id=call.data.get("sensor_id"),
@@ -173,10 +174,10 @@ async def send_frbc_instruction(
     ):  # pylint: disable=possibly-unused-variable
         """Send S2 Fill Rate Based Control message to the ResourceManager"""
 
-        if "cem" not in hass.data[DOMAIN]:
+        if "cem" not in hass.data[DOMAIN][entry.entry_id]:
             raise UndefinedCEMError()
 
-        cem: CEM = hass.data[DOMAIN]["cem"]
+        cem: CEM = hass.data[DOMAIN][entry.entry_id]["cem"]
 
         tz = pytz.timezone(hass.config.time_zone)
         DT_FMT = "%Y-%m-%d %H:%M:%S"
@@ -202,7 +203,7 @@ async def send_frbc_instruction(
     async def get_measurements(
         call: ServiceCall,
     ) -> ServiceResponse:  # pylint: disable=possibly-unused-variable
-        client: FlexMeasuresClient = hass.data[DOMAIN]["fm_client"]
+        client: FlexMeasuresClient = hass.data[DOMAIN][entry.entry_id]["fm_client"]
 
         response = await client.get_sensor_data(
             sensor_id=call.data.get("sensor_id"),

custom_components/flexmeasures_hacs/websockets.py

@@ -17,6 +17,7 @@
 from s2python.common import EnergyManagementRole, Handshake, ControlType
 
 from homeassistant.components.http import HomeAssistantView
+from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 
 from .const import DOMAIN, WS_VIEW_NAME, WS_VIEW_URI
@@ -32,7 +33,7 @@ class WebsocketAPIView(HomeAssistantView):
     url: str = WS_VIEW_URI
     requires_auth: bool = False
 
-    def __init__(self, entry) -> None:
+    def __init__(self, entry: ConfigEntry) -> None:
         """Initialize websocket view."""
         super().__init__()
         self.entry = entry
@@ -60,20 +61,22 @@ class WebSocketHandler:
 
     cem: CEM
 
-    def __init__(self, hass: HomeAssistant, entry, request: web.Request) -> None:
+    def __init__(
+        self, hass: HomeAssistant, entry: ConfigEntry, request: web.Request
+    ) -> None:
         """Initialize an active connection."""
         self.hass = hass
         self.request = request
         self.entry = entry
         self.wsock = web.WebSocketResponse(heartbeat=None)
 
         self.cem = CEM(
-            fm_client=hass.data[DOMAIN]["fm_client"],
+            fm_client=hass.data[DOMAIN][entry.entry_id]["fm_client"],
             default_control_type=ControlType.FILL_RATE_BASED_CONTROL,
         )
-        frbc_data: FRBC_Config = hass.data[DOMAIN]["frbc_config"]
+        frbc_data: FRBC_Config = hass.data[DOMAIN][entry.entry_id]["frbc_config"]
         frbc = FillRateBasedControlTUNES(**asdict(frbc_data))
-        hass.data[DOMAIN]["cem"] = self.cem
+        hass.data[DOMAIN][entry.entry_id]["cem"] = self.cem
         self.cem.register_control_type(frbc)
 
         self._logger = WebSocketAdapter(_WS_LOGGER, {"connid": id(self)})

tests/conftest.py

@@ -42,13 +42,15 @@ async def setup_fm_integration(hass: HomeAssistant):
             "soc_min": 0.0,
             "soc_max": 0.001,
         },
-        unique_id=1212121,
         state=ConfigEntryState.NOT_LOADED,
     )
 
     entry.add_to_hass(hass)
     assert await async_setup_component(hass, DOMAIN, {})
     await hass.async_block_till_done()
+    assert entry.entry_id in hass.data[DOMAIN]
+    print(f"ENTRY ID = {entry.entry_id}")
+    print(f"HASS DATA FOR OUR DOMAIN = {hass.data[DOMAIN]}")
 
     return entry