From 06a20d0d1540d35dc4e2446592e1a8b08ef56420 Mon Sep 17 00:00:00 2001
From: William Scanlon <wjs.scanlon@gmail.com>
Date: Tue, 22 Aug 2017 09:37:06 -0400
Subject: [PATCH] Fix octoprint errors when printer is off/disconnected (#8988)

* Fix octoprint errors when printer is off/disconnected
---
 homeassistant/components/octoprint.py        | 28 ++++++++++++++++----
 homeassistant/components/sensor/octoprint.py | 18 ++++++++++++-
 2 files changed, 40 insertions(+), 6 deletions(-)

diff --git a/homeassistant/components/octoprint.py b/homeassistant/components/octoprint.py
index 204490ce36c..fdf237d7180 100644
--- a/homeassistant/components/octoprint.py
+++ b/homeassistant/components/octoprint.py
@@ -16,11 +16,15 @@ import homeassistant.helpers.config_validation as cv
 _LOGGER = logging.getLogger(__name__)
 
 DOMAIN = 'octoprint'
+CONF_NUMBER_OF_TOOLS = 'number_of_tools'
+CONF_BED = 'bed'
 
 CONFIG_SCHEMA = vol.Schema({
     DOMAIN: vol.Schema({
         vol.Required(CONF_API_KEY): cv.string,
         vol.Required(CONF_HOST): cv.string,
+        vol.Optional(CONF_NUMBER_OF_TOOLS, default=0): cv.positive_int,
+        vol.Optional(CONF_BED, default=False): cv.boolean
     }),
 }, extra=vol.ALLOW_EXTRA)
 
@@ -29,11 +33,13 @@ def setup(hass, config):
     """Set up the OctoPrint component."""
     base_url = 'http://{}/api/'.format(config[DOMAIN][CONF_HOST])
     api_key = config[DOMAIN][CONF_API_KEY]
+    number_of_tools = config[DOMAIN][CONF_NUMBER_OF_TOOLS]
+    bed = config[DOMAIN][CONF_BED]
 
     hass.data[DOMAIN] = {"api": None}
 
     try:
-        octoprint_api = OctoPrintAPI(base_url, api_key)
+        octoprint_api = OctoPrintAPI(base_url, api_key, bed, number_of_tools)
         hass.data[DOMAIN]["api"] = octoprint_api
         octoprint_api.get('printer')
         octoprint_api.get('job')
@@ -46,7 +52,7 @@ def setup(hass, config):
 class OctoPrintAPI(object):
     """Simple JSON wrapper for OctoPrint's API."""
 
-    def __init__(self, api_url, key):
+    def __init__(self, api_url, key, bed, number_of_tools):
         """Initialize OctoPrint API and set headers needed later."""
         self.api_url = api_url
         self.headers = {'content-type': CONTENT_TYPE_JSON,
@@ -58,11 +64,23 @@ class OctoPrintAPI(object):
         self.available = False
         self.printer_error_logged = False
         self.job_error_logged = False
+        self.bed = bed
+        self.number_of_tools = number_of_tools
+        _LOGGER.error(str(bed) + " " + str(number_of_tools))
 
     def get_tools(self):
-        """Get the dynamic list of tools that temperature is monitored on."""
-        tools = self.printer_last_reading[0]['temperature']
-        return tools.keys()
+        """Get the list of tools that temperature is monitored on."""
+        tools = []
+        if self.number_of_tools > 0:
+            for tool_number in range(0, self.number_of_tools):
+                tools.append("tool" + str(tool_number))
+        if self.bed:
+            tools.append('bed')
+        if not self.bed and self.number_of_tools == 0:
+            temps = self.printer_last_reading[0].get('temperature')
+            if temps is not None:
+                tools = temps.keys()
+        return tools
 
     def get(self, endpoint):
         """Send a get request, and return the response as a dict."""
diff --git a/homeassistant/components/sensor/octoprint.py b/homeassistant/components/sensor/octoprint.py
index e3a06e58370..85b388a1919 100644
--- a/homeassistant/components/sensor/octoprint.py
+++ b/homeassistant/components/sensor/octoprint.py
@@ -20,6 +20,8 @@ _LOGGER = logging.getLogger(__name__)
 DEPENDENCIES = ['octoprint']
 DOMAIN = "octoprint"
 DEFAULT_NAME = 'OctoPrint'
+NOTIFICATION_ID = 'octoprint_notification'
+NOTIFICATION_TITLE = 'OctoPrint sensor setup error'
 
 SENSOR_TYPES = {
     'Temperatures': ['printer', 'temperature', '*', TEMP_CELSIUS],
@@ -42,12 +44,26 @@ def setup_platform(hass, config, add_devices, discovery_info=None):
     octoprint_api = hass.data[DOMAIN]["api"]
     name = config.get(CONF_NAME)
     monitored_conditions = config.get(CONF_MONITORED_CONDITIONS)
+    tools = octoprint_api.get_tools()
+    _LOGGER.error(str(tools))
+
+    if "Temperatures" in monitored_conditions:
+        if not tools:
+            hass.components.persistent_notification.create(
+                'Your printer appears to be offline.<br />'
+                'If you do not want to have your printer on <br />'
+                ' at all times, and you would like to monitor <br /> '
+                'temperatures, please add <br />'
+                'bed and/or number&#95of&#95tools to your config <br />'
+                'and restart.',
+                title=NOTIFICATION_TITLE,
+                notification_id=NOTIFICATION_ID)
 
     devices = []
     types = ["actual", "target"]
     for octo_type in monitored_conditions:
         if octo_type == "Temperatures":
-            for tool in octoprint_api.get_tools():
+            for tool in tools:
                 for temp_type in types:
                     new_sensor = OctoPrintSensor(
                         octoprint_api, temp_type, temp_type, name,