From 332b8b0341876f8df4073e924a59ef70c5c6558e Mon Sep 17 00:00:00 2001
From: puddly <32534428+puddly@users.noreply.github.com>
Date: Wed, 13 Jul 2022 08:36:48 -0400
Subject: [PATCH] Load group information from existing DB device

---
 bellows/zigbee/application.py |  9 +++++++
 tests/test_application.py     | 44 +++++++++++++++++++++++++++++++++++
 2 files changed, 53 insertions(+)

diff --git a/bellows/zigbee/application.py b/bellows/zigbee/application.py
index 384fcffe..64135b07 100644
--- a/bellows/zigbee/application.py
+++ b/bellows/zigbee/application.py
@@ -185,6 +185,11 @@ async def start_network(self):
         self.controller_event.set()
         self._watchdog_task = asyncio.create_task(self._watchdog())
 
+        try:
+            db_device = self.get_device(ieee=self.state.node_info.ieee)
+        except KeyError:
+            db_device = None
+
         ezsp_device = EZSPCoordinator(
             application=self,
             ieee=self.state.node_info.ieee,
@@ -198,6 +203,10 @@ async def start_network(self):
         ezsp_device.manufacturer = ezsp_device.endpoints[1].manufacturer
         await ezsp_device.schedule_initialize()
 
+        # Group membership is stored in the database for EZSP coordinators
+        if db_device is not None:
+            ezsp_device.endpoints[1].member_of.update(db_device.endpoints[1].member_of)
+
         await self.multicast.startup(ezsp_device)
 
     async def load_network_info(self, *, load_devices=False) -> None:
diff --git a/tests/test_application.py b/tests/test_application.py
index 677f371d..1682d5d0 100644
--- a/tests/test_application.py
+++ b/tests/test_application.py
@@ -1539,3 +1539,47 @@ async def test_ensure_network_running_not_joined_success(app):
 
     ezsp.networkState.assert_called_once()
     ezsp.networkInit.assert_called_once()
+
+
+async def test_startup_coordinator_existing_groups_joined(app, ieee):
+    """Coordinator joins groups loaded from the database."""
+
+    app._ensure_network_running = AsyncMock()
+    app._ezsp.update_policies = AsyncMock()
+    app.load_network_info = AsyncMock()
+
+    app._multicast = bellows.multicast.Multicast(app._ezsp)
+    app.state.node_info.ieee = ieee
+
+    db_device = app.add_device(ieee, 0x0000)
+    db_ep = db_device.add_endpoint(1)
+
+    app.groups.add_group(0x1234, "Group Name", suppress_event=True)
+    app.groups[0x1234].add_member(db_ep, suppress_event=True)
+
+    p1 = patch.object(bellows.multicast.Multicast, "_initialize")
+    p2 = patch.object(bellows.multicast.Multicast, "subscribe")
+
+    with p1 as p1, p2 as p2:
+        await app.start_network()
+
+    p2.assert_called_once_with(0x1234)
+
+
+async def test_startup_new_coordinator_no_groups_joined(app, ieee):
+    """Coordinator freshy added to the database has no groups to join."""
+
+    app._ensure_network_running = AsyncMock()
+    app._ezsp.update_policies = AsyncMock()
+    app.load_network_info = AsyncMock()
+
+    app._multicast = bellows.multicast.Multicast(app._ezsp)
+    app.state.node_info.ieee = ieee
+
+    p1 = patch.object(bellows.multicast.Multicast, "_initialize")
+    p2 = patch.object(bellows.multicast.Multicast, "subscribe")
+
+    with p1 as p1, p2 as p2:
+        await app.start_network()
+
+    p2.assert_not_called()