From febc00b6db86d8f898397408eebcf7a71af4947e Mon Sep 17 00:00:00 2001
From: Jacob Mansfield <cyberjacob@gmail.com>
Date: Tue, 4 Sep 2018 20:15:58 +0100
Subject: [PATCH] Automatic reconnect (#4)

* Update asterisk_ami.py

* Update asterisk_ami.py
---
 asterisk_ami.py        | 32 +++++++++++++++++++-------------
 sensor/asterisk_ami.py | 12 ++++++------
 2 files changed, 25 insertions(+), 19 deletions(-)

diff --git a/asterisk_ami.py b/asterisk_ami.py
index 24a1b8a..4b91f9d 100644
--- a/asterisk_ami.py
+++ b/asterisk_ami.py
@@ -40,21 +40,9 @@ def setup(hass, config):
     import asterisk.manager
     manager = asterisk.manager.Manager()
 
-    host = config[DOMAIN].get(CONF_HOST, DEFAULT_HOST)
-    port = config[DOMAIN].get(CONF_PORT, DEFAULT_PORT)
-    username = config[DOMAIN].get(CONF_USERNAME, DEFAULT_USERNAME)
-    password = config[DOMAIN].get(CONF_PASSWORD, DEFAULT_PASSWORD)
-
-    try:
-        manager.connect(host, port)
-        login_status = manager.login(username=username, secret=password).get_header("Response")
-    except asterisk.manager.ManagerException as exception:
-        _LOGGER.error("Error connecting to Asterisk: %s", exception.args[1])
+    if not connect(hass):
         return False
 
-    if "Success" not in login_status:
-        _LOGGER.error("Could not authenticate: %s", login_status)
-
     hass.data[DATA_ASTERISK] = manager
     hass.data[DATA_MONITOR] = config[DOMAIN].get(CONF_MONITOR, [])
     hass.data[DATA_MAILBOX] = config[DOMAIN].get(CONF_MAILBOX, [])
@@ -62,3 +50,21 @@ def setup(hass, config):
     load_platform(hass, 'sensor', DOMAIN)
 
     return True
+
+def connect(hass):
+    host = hass.config[DOMAIN].get(CONF_HOST, DEFAULT_HOST)
+    port = hass.config[DOMAIN].get(CONF_PORT, DEFAULT_PORT)
+    username = hass.config[DOMAIN].get(CONF_USERNAME, DEFAULT_USERNAME)
+    password = hass.config[DOMAIN].get(CONF_PASSWORD, DEFAULT_PASSWORD)
+
+    try:
+        hass.data[DATA_ASTERISK].connect(host, port)
+        login_status = hass.data[DATA_ASTERISK].login(username=username, secret=password).get_header("Response")
+    except asterisk.manager.ManagerException as exception:
+        _LOGGER.error("Error connecting to Asterisk: %s", exception.args[1])
+        return False
+    if "Success" not in login_status:
+        _LOGGER.error("Could not authenticate: %s", login_status)
+        return False
+    _LOGGER.info("Successfully connected to Asterisk server")
+    return True
diff --git a/sensor/asterisk_ami.py b/sensor/asterisk_ami.py
index dbf0b3d..523673f 100644
--- a/sensor/asterisk_ami.py
+++ b/sensor/asterisk_ami.py
@@ -4,9 +4,7 @@
 from homeassistant.const import STATE_UNKNOWN
 from homeassistant.helpers.entity import Entity
 
-DATA_ASTERISK = 'asterisk'
-DATA_MONITOR = 'asterisk-monitor'
-DATA_MAILBOX = 'asterisk-mailbox'
+from .. import asterisk_ami as ami_platform
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -15,9 +13,9 @@ def setup_platform(hass, config, add_devices, discovery_info=None):
        Loads configuration and creates devices for extensions and mailboxes
     """
     add_devices([AsteriskSensor(hass)])
-    for extension in hass.data[DATA_MONITOR]:
+    for extension in hass.data[ami_platform.DATA_MONITOR]:
         add_devices([AsteriskExtension(hass, extension)])
-    for mailbox in hass.data[DATA_MAILBOX]:
+    for mailbox in hass.data[ami_platform.DATA_MAILBOX]:
         add_devices([AsteriskMailbox(hass, mailbox)])
 
 
@@ -45,11 +43,13 @@ def should_poll(self):
 
     def update(self):
         """Check the connection status and update the internal state"""
-        if self._hass.data[DATA_ASTERISK].connected():
+        if self._hass.data[ami_platform.DATA_ASTERISK].connected():
             self._state = 'connected'
         else:
             self._state = 'disconnected'
 
+            ami_platform.connect(self._hass)
+
 
 class AsteriskExtension(Entity):
     """SIP extension connection status (registration) sensor."""