v0.1.0 - automatic reconnect, request assignment

.gitignore

@@ -103,4 +103,5 @@ venv.bak/
 # mypy
 .mypy_cache/
 
-venv/
+venv/
+hass_debug_env/

.vscode/launch.json

@@ -0,0 +1,23 @@
+{
+    // Use IntelliSense to learn about possible attributes.
+    // Hover to view descriptions of existing attributes.
+    // For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387
+    "version": "0.2.0",
+    "configurations": [
+        {
+            "name": "Python: hass",
+            "type": "python",
+            "request": "launch",
+            "justMyCode": false,
+            "module": "homeassistant",
+            "args": ["-c", "./hass_debug_env"]
+        },
+        {
+            "name": "Python: Current File",
+            "type": "python",
+            "request": "launch",
+            "program": "${file}",
+            "console": "integratedTerminal"
+        }
+    ]
+}

.vscode/settings.json

@@ -0,0 +1,3 @@
+{
+    "python.pythonPath": "venv/bin/python"
+}

custom_components/foldingathomecontrol/__init__.py

@@ -4,8 +4,7 @@
 
 import asyncio
 import itertools
-import logging
-from typing import Any
+from typing import Any, Tuple
 
 import homeassistant.helpers.config_validation as cv
 import voluptuous as vol
@@ -26,9 +25,10 @@
     DOMAIN,
     SENSOR_ADDED,
     SENSOR_REMOVED,
+    _LOGGER,
 )
 
-_LOGGER = logging.getLogger(__name__)
+from .services import async_setup_services, async_unload_services
 
 FOLDINGATHOMECONTROL_SCHEMA = vol.Schema(
     {
@@ -64,13 +64,19 @@ async def async_setup_entry(hass, config_entry):
     if not await data.async_setup():
         return False
 
+    await async_setup_services(hass)
+
     return True
 
 
 async def async_unload_entry(hass, config_entry):
     """Unload a config entry."""
     await hass.config_entries.async_forward_entry_unload(config_entry, "sensor")
+    await hass.data[DOMAIN][config_entry.entry_id].async_remove()
     hass.data[DOMAIN].pop(config_entry.entry_id)
+    # If there is no instance of this integration registered anymore
+    if not hass.data[DOMAIN]:
+        await async_unload_services(hass)
     return True
 
 
@@ -87,30 +93,12 @@ def __init__(self, hass: HomeAssistant, config_entry) -> None:
         self._task = None
 
     @callback
-    def callback(self, message_type: str, data: Any) -> None:
+    def data_received_callback(self, message_type: str, data: Any) -> None:
         """Called when data is received from the Folding@Home client."""
         if message_type == PyOnMessageTypes.SLOTS.value:
-            if PyOnMessageTypes.SLOTS.value in self.data:
-                added = list(
-                    itertools.filterfalse(
-                        lambda x: x in self.data[PyOnMessageTypes.SLOTS.value], data
-                    )
-                )
-                removed = list(
-                    itertools.filterfalse(
-                        lambda x: x in data, self.data[PyOnMessageTypes.SLOTS.value]
-                    )
-                )
-                async_dispatcher_send(
-                    self.hass, self.get_sensor_added_identifer(), added
-                )
-                async_dispatcher_send(
-                    self.hass, self.get_sensor_removed_identifer(), removed
-                )
-            else:
-                async_dispatcher_send(
-                    self.hass, self.get_sensor_added_identifer(), data
-                )
+            self.handle_slots_data_received(data)
+        if message_type == PyOnMessageTypes.ERROR.value:
+            self.handle_error_received(data)
         self.data[message_type] = data
         async_dispatcher_send(self.hass, self.get_data_update_identifer())
 
@@ -122,6 +110,7 @@ async def async_setup(self) -> bool:
         self.client = FoldingAtHomeController(address, port, password)
         try:
             await self.client.try_connect_async(timeout=5)
+            await self.client.subscribe_async()
         except FoldingAtHomeControlConnectionFailed:
             return False
 
@@ -131,8 +120,10 @@ async def async_setup(self) -> bool:
             )
         )
 
-        self._remove_callback = self.client.register_callback(self.callback)
-        self._task = asyncio.ensure_future(self.client.run())
+        self._remove_callback = self.client.register_callback(
+            self.data_received_callback
+        )
+        self._task = asyncio.ensure_future(self.client.start())
 
         return True
 
@@ -142,6 +133,7 @@ async def async_remove(self) -> None:
             self._remove_callback()
         if self._task is not None:
             self._task.cancel()
+            await self._task
 
     def get_data_update_identifer(self) -> None:
         """Returns the unique data_update itentifier for this connection."""
@@ -155,6 +147,42 @@ def get_sensor_removed_identifer(self) -> None:
         """Returns the unique sensor_removed itentifier for this connection."""
         return f"{SENSOR_REMOVED}_{self.config_entry.data[CONF_ADDRESS]}"
 
+    def handle_error_received(self, error: Any) -> None:
+        """Handle received error message."""
+        _LOGGER.warning(
+            "%s received error: %s", self.config_entry.data[CONF_ADDRESS], error
+        )
+
+    def handle_slots_data_received(self, data: Any) -> None:
+        """Handle received slots data."""
+        if PyOnMessageTypes.SLOTS.value in self.data:
+            added, removed = self.calculate_slot_changes(data)
+            async_dispatcher_send(self.hass, self.get_sensor_added_identifer(), added)
+            async_dispatcher_send(
+                self.hass, self.get_sensor_removed_identifer(), removed
+            )
+        else:
+            async_dispatcher_send(self.hass, self.get_sensor_added_identifer(), data)
+
+    def calculate_slot_changes(self, slots: dict) -> Tuple[dict, dict]:
+        """Get added and removed slots."""
+        added = list(
+            itertools.filterfalse(
+                lambda slot: (
+                    slot["id"] != entry["id"]
+                    for entry in self.data[PyOnMessageTypes.SLOTS.value]
+                ),
+                slots,
+            )
+        )
+        removed = list(
+            itertools.filterfalse(
+                lambda entry: (entry["id"] != slot["id"] for slot in slots),
+                self.data[PyOnMessageTypes.SLOTS.value],
+            )
+        )
+        return added, removed
+
     @property
     def available(self):
         """Is the Folding@Home client available."""

custom_components/foldingathomecontrol/const.py

@@ -1,13 +1,18 @@
 """Constants for foldingathomecontrol."""
+import logging
+
 # Base component constants
 DOMAIN = "foldingathomecontrol"
 DOMAIN_DATA = f"{DOMAIN}_data"
-VERSION = "0.0.1"
+VERSION = "0.1.0"
 PLATFORMS = ["binary_sensor", "sensor"]
 DATA_UPDATED = f"{DOMAIN}_data_updated"
 SENSOR_ADDED = f"{DOMAIN}_sensor_added"
 SENSOR_REMOVED = f"{DOMAIN}_sensor_removed"
 
+# Logger
+_LOGGER = logging.getLogger(__package__)
+
 # Icons
 ICON = "mdi:state-machine"
 

custom_components/foldingathomecontrol/manifest.json

@@ -8,7 +8,7 @@
     "@eifinger"
   ],
   "requirements": [
-    "PyFoldingAtHomeControl==0.1.6"
+    "PyFoldingAtHomeControl==1.0.0"
   ],
   "homeassistant": "0.96.0"
 }

custom_components/foldingathomecontrol/services.py

@@ -0,0 +1,66 @@
+"""FoldingAtHomeControl services."""
+import voluptuous as vol
+
+from homeassistant.helpers import config_validation as cv
+
+from .const import (
+    CONF_ADDRESS,
+    _LOGGER,
+    DOMAIN,
+)
+
+DOMAIN_SERVICES = f"{DOMAIN}_services"
+
+SERVICE_ADDRESS = "address"
+
+SERVICE_REQUEST_WORK_SERVER_ASSIGNMENT = "request_work_server_assignment"
+SERVICE_REQUEST_WORK_SERVER_ASSIGNMENT_SCHEMA = vol.Schema(
+    {vol.Required(SERVICE_ADDRESS): cv.string}
+)
+
+
+async def async_setup_services(hass):
+    """Set up services for FoldingAtHomeControl integration."""
+    if hass.data.get(DOMAIN_SERVICES, False):
+        return
+
+    hass.data[DOMAIN_SERVICES] = True
+
+    async def async_call_folding_at_home_control_service(service_call):
+        """Call correct FoldingAtHomeControl service."""
+        service = service_call.service
+        service_data = service_call.data
+
+        if service == SERVICE_REQUEST_WORK_SERVER_ASSIGNMENT:
+            await async_request_assignment_service(hass, service_data)
+
+    hass.services.async_register(
+        DOMAIN,
+        SERVICE_REQUEST_WORK_SERVER_ASSIGNMENT,
+        async_call_folding_at_home_control_service,
+        schema=SERVICE_REQUEST_WORK_SERVER_ASSIGNMENT_SCHEMA,
+    )
+
+
+async def async_unload_services(hass):
+    """Unload deCONZ services."""
+    if not hass.data.get(DOMAIN_SERVICES):
+        return
+
+    hass.data[DOMAIN_SERVICES] = False
+
+    hass.services.async_remove(DOMAIN, SERVICE_REQUEST_WORK_SERVER_ASSIGNMENT)
+
+
+async def async_request_assignment_service(hass, data):
+    """Let the client request a new work server assignment."""
+
+    address = data[SERVICE_ADDRESS]
+
+    for config_entry in hass.data[DOMAIN]:
+        if hass.data[DOMAIN][config_entry].config_entry.data[CONF_ADDRESS] == address:
+            await hass.data[DOMAIN][
+                config_entry
+            ].client.request_work_server_assignment()
+            return
+    _LOGGER.warning("Could not find a registered integration with address: %s", address)

custom_components/foldingathomecontrol/services.yaml

@@ -0,0 +1,7 @@
+request_work_server_assignment:
+    description: Request a new assignment from the work server.
+    fields:
+      address:
+        description: The IP address or hostname of the client. It can be found as part of the integration name.
+        example: 'localhost'
+

requirements.txt

@@ -1 +1 @@
-PyFoldingAtHomeControl==0.1.6
+PyFoldingAtHomeControl==1.0.0