From 46d71769765ba744bf20f5bd63aefc034bac769c Mon Sep 17 00:00:00 2001
From: Kryzek <dr.drake@gmail.com>
Date: Sat, 23 Nov 2019 12:47:37 +0200
Subject: [PATCH] availability_whitelist (#2387)

* availability_whitelist

* Tests

* Cleanup

* Update deviceAvailability.js
---
 lib/extension/deviceAvailability.js | 12 ++++++++++++
 lib/util/settings.js                |  2 ++
 test/deviceAvailability.test.js     | 29 +++++++++++++++++++++++++++++
 3 files changed, 43 insertions(+)

diff --git a/lib/extension/deviceAvailability.js b/lib/extension/deviceAvailability.js
index 98e8c7f8..38aa0f13 100644
--- a/lib/extension/deviceAvailability.js
+++ b/lib/extension/deviceAvailability.js
@@ -27,13 +27,25 @@ class DeviceAvailability extends BaseExtension {
         this.blacklist = settings.get().advanced.availability_blacklist.map((e) => {
             return settings.getEntity(e).ID;
         });
+
+        // Initialize whitelist
+        this.whitelist = settings.get().advanced.availability_whitelist.map((e) => {
+            return settings.getEntity(e).ID;
+        });
     }
 
     isPingable(device) {
+        // Whitelist is not empty and device is in it, enable availability
+        if (this.whitelist.length > 0) {
+            return this.whitelist.includes(device.ieeeAddr);
+        }
+
+        // Device is on blacklist, disable availability
         if (this.blacklist.includes(device.ieeeAddr)) {
             return false;
         }
 
+        // Device is on forcedPingable-list, enable availability
         if (forcedPingable.find((d) => d.zigbeeModel.includes(device.modelID))) {
             return true;
         }
diff --git a/lib/util/settings.js b/lib/util/settings.js
index 79a0426e..6ff635cf 100644
--- a/lib/util/settings.js
+++ b/lib/util/settings.js
@@ -54,6 +54,7 @@ const defaults = {
         // Availability timeout in seconds, disabled by default.
         availability_timeout: 0,
         availability_blacklist: [],
+        availability_whitelist: [],
 
         /**
          * Home Assistant requires ALL attributes to be present in ALL MQTT messages send by the device.
@@ -155,6 +156,7 @@ const schema = {
                 elapsed: {type: 'boolean'},
                 availability_timeout: {type: 'number', minimum: 0},
                 availability_blacklist: {type: 'array', items: {type: 'string'}},
+                availability_whitelist: {type: 'array', items: {type: 'string'}},
                 report: {type: 'boolean'},
                 homeassistant_discovery_topic: {type: 'string'},
                 homeassistant_status_topic: {type: 'string'},
diff --git a/test/deviceAvailability.test.js b/test/deviceAvailability.test.js
index 105bc82e..01cebae4 100644
--- a/test/deviceAvailability.test.js
+++ b/test/deviceAvailability.test.js
@@ -231,6 +231,35 @@ describe('Device availability', () => {
         expect(device.ping).toHaveBeenCalledTimes(1);
     });
 
+    it('Should ping whitelisted devices if availability_whitelist is set', async () => {
+        const device = zigbeeHerdsman.devices.bulb_color;
+        settings.set(['advanced', 'availability_whitelist'], [device.ieeeAddr])
+        await controller.stop();
+        await flushPromises();
+        controller = new Controller();
+        await controller.start();
+        await flushPromises();
+        device.ping.mockClear();
+        jest.advanceTimersByTime(11 * 1000);
+        await flushPromises();
+        expect(device.ping).toHaveBeenCalledTimes(1);
+    });
+
+    it('Should not ping non-whitelisted devices if availability_whitelist is set', async () => {
+        const device = zigbeeHerdsman.devices.bulb;
+        getExtension().state[device.ieeeAddr] = false;
+        settings.set(['advanced', 'availability_whitelist'], [device.ieeeAddr])
+        await controller.stop();
+        await flushPromises();
+        controller = new Controller();
+        await controller.start();
+        await flushPromises();
+        device.ping.mockClear();
+        jest.advanceTimersByTime(11 * 1000);
+        await flushPromises();
+        expect(device.ping).toHaveBeenCalledTimes(0);
+    });    
+
     it('Should not read when device has no modelID and reconnects', async () => {
         const device = zigbeeHerdsman.devices.nomodel;
         getExtension().state[device.ieeeAddr] = true;