From 500bce2418aa026a268523940f754685ecb0b40b Mon Sep 17 00:00:00 2001
From: Koen Kanters <koenkanters94@gmail.com>
Date: Mon, 23 Apr 2018 18:17:47 +0200
Subject: [PATCH] Perform home assistant MQTT discovery on startup.

---
 lib/controller.js    | 26 ++++++++++++++++++++------
 lib/mqtt.js          |  8 ++++++--
 lib/util/settings.js |  1 +
 lib/zigbee.js        |  6 +++++-
 4 files changed, 32 insertions(+), 9 deletions(-)

diff --git a/lib/controller.js b/lib/controller.js
index a8208e71f..c780c6526 100644
--- a/lib/controller.js
+++ b/lib/controller.js
@@ -23,7 +23,21 @@ class Controller {
             if (error) {
                 logger.error('Failed to start');
             } else {
-                this.mqtt.connect(this.handleMQTTMessage);
+                this.mqtt.connect(this.handleMQTTMessage, () => {
+                    // Home assistant MQTT discovery on startup.
+                    if (settings.get().homeassistant_discovery) {
+                        const devices = this.zigbee.getAllClients();
+
+                        devices.forEach((device) => {
+                            const mappedModel = deviceMapping[device.modelId];
+                            
+                            if (mappedModel && mappedModel.homeassistant && !this.hassDiscoveryCache[device.ieeeAddr]) {
+                                this.homeassistantDiscover(mappedModel, device, settings.getDevice(device.ieeeAddr).friendly_name);
+                                this.hassDiscoveryCache[device.ieeeAddr] = true;
+                            }
+                        });
+                    }
+                });
             }
         });
     }
@@ -49,7 +63,7 @@ class Controller {
         const device = message.endpoints[0].device;
 
         // Check if this is a new device.
-        if (!settings.get().devices[device.ieeeAddr]) {
+        if (!settings.getDevice(device.ieeeAddr)) {
             logger.info(`New device with address ${device.ieeeAddr} connected!`);
 
             settings.get().devices[device.ieeeAddr] = {
@@ -68,7 +82,7 @@ class Controller {
         // Map Zigbee modelID to vendor modelID.
         const modelID = message.endpoints[0].device.modelId;
         const mappedModel = deviceMapping[modelID];
-        const friendlyName = settings.get().devices[device.ieeeAddr].friendly_name;
+        const friendlyName = settings.getDevice(device.ieeeAddr).friendly_name;
 
         if (!mappedModel) {
             logger.warn(`Device with modelID '${modelID}' is not supported.`);
@@ -99,7 +113,7 @@ class Controller {
         }
 
         // Convert this Zigbee message to a MQTT message.
-        const retain = settings.get().devices[device.ieeeAddr].retain;
+        const retain = settings.getDevice(device.ieeeAddr).retain;
 
         const publish = (payload) => {
             if (this.stateCache[device.ieeeAddr]) {
@@ -130,7 +144,7 @@ class Controller {
         const friendlyName = topic.split('/')[1];
 
         // Map friendlyName to deviceID.
-        const deviceID = Object.keys(settings.get().devices).find((id) => settings.get().devices[id].friendly_name === friendlyName);
+        const deviceID = Object.keys(settings.get().devices).find((id) => settings.getDevice(id).friendly_name === friendlyName);
         if (!deviceID) {
             logger.error(`Cannot handle '${topic}' because deviceID of '${friendlyName}' cannot be found`);
             return;
@@ -153,7 +167,7 @@ class Controller {
                     this.mqtt.publish(
                         friendlyName, 
                         JSON.stringify({state: json[key]}), 
-                        settings.get().devices[deviceID].retain,
+                        settings.getDevice(deviceID).retain,
                     );
                 }
             };
diff --git a/lib/mqtt.js b/lib/mqtt.js
index 955db3559..c89ccfdb5 100644
--- a/lib/mqtt.js
+++ b/lib/mqtt.js
@@ -10,7 +10,7 @@ class MQTT {
         this.handleMessage = this.handleMessage.bind(this);
     }
 
-    connect(onMessage) {
+    connect(onMessage, callback) {
         const mqttSettings = settings.get().mqtt;
         logger.info(`Connecting to MQTT server at ${mqttSettings.server}`);
 
@@ -23,7 +23,11 @@ class MQTT {
         this.client = mqtt.connect(mqttSettings.server, options);
 
         // Register callbacks.
-        this.client.on('connect', this.handleConnect);
+        this.client.on('connect', () => {
+            this.handleConnect();
+            callback();
+        });
+        
         this.client.on('message', this.handleMessage);
 
         // Set timer at interval to check if connected to MQTT server.
diff --git a/lib/util/settings.js b/lib/util/settings.js
index a40580717..ee8925062 100644
--- a/lib/util/settings.js
+++ b/lib/util/settings.js
@@ -21,4 +21,5 @@ function read() {
 module.exports = {
     get: () => settings,
     write: () => write(),
+    getDevice: (id) => settings.devices[id],
 }
diff --git a/lib/zigbee.js b/lib/zigbee.js
index e789374de..541290ff5 100644
--- a/lib/zigbee.js
+++ b/lib/zigbee.js
@@ -47,7 +47,7 @@ class Zigbee {
     handleReady() {
         logger.info('zigbee-shepherd ready');
 
-        const devices = this.shepherd.list().filter((device) => device.type !== 'Coordinator');
+        const devices = this.getAllClients();
 
         logger.info(`Currently ${devices.length} devices are joined:`);
         devices.forEach((device) => logger.info(this.getDeviceLogMessage(device)));
@@ -77,6 +77,10 @@ class Zigbee {
         });
     }
 
+    getAllClients() {
+        return this.shepherd.list().filter((device) => device.type !== 'Coordinator');
+    }
+
     handleMessage(message) {
         if (this.onMessage) {
             this.onMessage(message);