Isolate SSDP into separate module

Author: codingjoe

setup.cfg

@@ -23,5 +23,9 @@ keywords =
 packages =
     yeelib
 
+[pycodestyle]
+max-line-length=99
+exclude = fixtures.py
+
 [pydocstyle]
-add_ignore = D1
+add_ignore = D1,D401

tests/fixtures.py

@@ -0,0 +1,44 @@
+request = b"""NOTIFY * HTTP/1.1
+Host: 239.255.255.250:1982
+Cache-Control: max-age=3600
+Location: yeelight://192.168.1.239:55443
+NTS: ssdp:alive
+Server: POSIX, UPnP/1.0 YGLC/1
+id: 0x000000000015243f
+model: color
+fw_ver: 18
+support: get_prop set_default set_power toggle set_bright start_cf stop_cf set_scene cron_add cron_get cron_del set_ct_abx set_rgb
+power: on
+bright: 100
+color_mode: 2
+ct: 4000
+rgb: 16711680
+hue: 100
+sat: 35
+name: my_bulb""".replace(b'\n', b'\r\n')
+
+response = b"""HTTP/1.1 200 OK
+Cache-Control: max-age=3600
+Date:
+Ext:
+Location: yeelight://192.168.1.239:55443
+Server: POSIX UPnP/1.0 YGLC/1
+id: 0x000000000015243f
+model: color
+fw_ver: 18
+support: get_prop set_default set_power toggle set_bright start_cf stop_cf set_scene cron_add cron_get cron_del set_ct_abx set_rgb
+power: on
+bright: 100
+color_mode: 2
+ct: 4000
+rgb: 16711680
+hue: 100
+sat: 35
+name: my_bulb""".replace(b'\n', b'\r\n')
+
+response_wrong_location = b"""HTTP/1.1 200 OK
+Cache-Control: max-age=3600
+Date:
+Ext:
+Location: yeelight://not.an.ip:55443
+Server: POSIX UPnP/1.0 YGLC/1""".replace(b'\n', b'\r\n')

tests/test_bulbs.py

@@ -26,8 +26,9 @@ def test_kwargs(self):
             'id': '0x000000000015243f',
             'model': 'color',
             'fw_ver': 18,
-            'support': 'get_prop set_default set_power toggle set_bright start_cf stop_cf set_scene'
-                       'cron_add cron_get cron_del set_ct_abx set_rgb',
+            'support': 'get_prop set_default set_power toggle set_bright'
+                       ' start_cf stop_cf set_scene cron_add cron_get'
+                       ' cron_del set_ct_abx set_rgb',
             'power': 'on',
             'bright': 100,
             'color_mode': 2,
@@ -41,5 +42,5 @@ def test_kwargs(self):
         with Bulb(*self.bulb_addr, **kwargs) as b:
             assert b.fw_ver == 18
             assert b.support == ['get_prop', 'set_default', 'set_power', 'toggle', 'set_bright',
-                                 'start_cf', 'stop_cf', 'set_scenecron_add', 'cron_get',
+                                 'start_cf', 'stop_cf', 'set_scene', 'cron_add', 'cron_get',
                                  'cron_del', 'set_ct_abx', 'set_rgb']

tests/test_discover.py

@@ -5,83 +5,38 @@
 from yeelib.discover import YeelightProtocol, search_bulbs
 from yeelib.exceptions import YeelightError
 
-
-notify = b"""NOTIFY * HTTP/1.1
-Host: 239.255.255.250:1982
-Cache-Control: max-age=3600
-Location: yeelight://192.168.1.239:55443
-NTS: ssdp:alive
-Server: POSIX, UPnP/1.0 YGLC/1
-id: 0x000000000015243f
-model: color
-fw_ver: 18
-support: get_prop set_default set_power toggle set_bright start_cf stop_cf set_scene
-cron_add cron_get cron_del set_ct_abx set_rgb
-power: on
-bright: 100
-color_mode: 2
-ct: 4000
-rgb: 16711680
-hue: 100
-sat: 35
-name: my_bulb"""
-
-mcast = b"""HTTP/1.1 200 OK
-Cache-Control: max-age=3600
-Date:
-Ext:
-Location: yeelight://192.168.1.239:55443
-Server: POSIX UPnP/1.0 YGLC/1
-id: 0x000000000015243f
-model: color
-fw_ver: 18
-support: get_prop set_default set_power toggle set_bright start_cf stop_cf set_scene
-cron_add cron_get cron_del set_ct_abx set_rgb
-power: on
-bright: 100
-color_mode: 2
-ct: 4000
-rgb: 16711680
-hue: 100
-sat: 35
-name: my_bulb"""
-
-wrong_location = b"""HTTP/1.1 200 OK
-Cache-Control: max-age=3600
-Date:
-Ext:
-Location: yeelight://not.an.ip:55443
-Server: POSIX UPnP/1.0 YGLC/1"""
+from . import fixtures
 
 
 class TestYeelightProtocoll:
     def test_notify(self, ):
         bulbs = {}
         p = YeelightProtocol(bulbs=bulbs)
-        p.datagram_received(data=notify, addr=('192.168.1.239', 1982))
+        p.datagram_received(data=fixtures.request, addr=('192.168.1.239', 1982))
         assert len(bulbs) == 1
         assert bulbs['0x000000000015243f'].ip == '192.168.1.239'
 
     def test_mcast(self, ):
         bulbs = {}
         p = YeelightProtocol(bulbs=bulbs)
-        p.datagram_received(data=mcast, addr=('192.168.1.239', 1982))
+        p.datagram_received(data=fixtures.response, addr=('192.168.1.239', 1982))
         assert len(bulbs) == 1
         assert bulbs['0x000000000015243f'].ip == '192.168.1.239'
 
     def test_duplicate(self):
         bulbs = {}
         p = YeelightProtocol(bulbs=bulbs)
-        p.datagram_received(data=notify, addr=('192.168.1.239', 1982))
-        p.datagram_received(data=notify, addr=('192.168.1.239', 1982))
+        p.datagram_received(data=fixtures.request, addr=('192.168.1.239', 1982))
+        p.datagram_received(data=fixtures.request, addr=('192.168.1.239', 1982))
         assert len(bulbs) == 1
         assert bulbs['0x000000000015243f'].ip == '192.168.1.239'
 
     def test_wrong_location(self):
         bulbs = {}
         p = YeelightProtocol(bulbs=bulbs)
         with pytest.raises(YeelightError) as e:
-            p.datagram_received(data=wrong_location, addr=('192.168.1.239', 1982))
+            p.datagram_received(data=fixtures.response_wrong_location,
+                                addr=('192.168.1.239', 1982))
         assert 'Location does not match: yeelight://not.an.ip:55443' in str(e)
 
 

tests/test_upnp.py

@@ -0,0 +1,54 @@
+import pytest
+
+from yeelib.upnp import SSDPMessage, SSDPResponse, SSDPRequest
+from . import fixtures
+
+
+class TestSSDPMessage:
+    def test_headers_copy(self):
+        headers = [('Cache-Control', 'max-age=3600')]
+        msg = SSDPMessage(headers=headers)
+        assert msg.headers == headers
+        assert msg.headers is not headers
+
+    def test_headers_dict(self):
+        headers = {'Cache-Control': 'max-age=3600'}
+        msg = SSDPMessage(headers=headers)
+        assert msg.headers == [('Cache-Control', 'max-age=3600')]
+
+    def test_headers_none(self):
+        msg = SSDPMessage(headers=None)
+        assert msg.headers == []
+
+    def test_parse(self):
+        with pytest.raises(NotImplementedError):
+            SSDPMessage.parse('')
+
+    def test_parse_headers(self):
+        headers = SSDPMessage.parse_headers('Cache-Control: max-age=3600')
+        assert headers == [('Cache-Control', 'max-age=3600')]
+
+
+class TestSSDPResponse:
+    def test_parse(self):
+        response = SSDPResponse.parse(fixtures.response.decode())
+        assert response.status_code == 200
+        assert response.reason == 'OK'
+
+
+class TestSSDPRequest:
+    def test_parse(self):
+        request = SSDPRequest.parse(fixtures.request.decode())
+        assert request.method == 'NOTIFY'
+        assert request.uri == '*'
+
+    def test_str(self):
+        request = SSDPRequest('NOTIFY', '*', headers=[('Cache-Control', 'max-age=3600')])
+        assert str(request) == (
+            'NOTIFY * HTTP/1.1\n'
+            'Cache-Control: max-age=3600'
+        )
+
+    def test_bytes(self):
+        request = SSDPRequest.parse(fixtures.request.decode())
+        assert bytes(request) == fixtures.request

yeelib/discover.py

@@ -1,6 +1,4 @@
 import asyncio
-import email.parser
-import errno
 import fcntl
 import logging
 import os
@@ -11,21 +9,15 @@
 from contextlib import contextmanager
 
 from yeelib.bulbs import Bulb
-
+from yeelib.upnp import SSDPRequest, SimpleServiceDiscoveryProtocol
 from .exceptions import YeelightError
 
-__all__ = ('search_bulbs',)
+__all__ = ('search_bulbs', 'YeelightProtocol')
 
 logger = logging.getLogger('yeelib')
 
-MCAST_IP = '239.255.255.250'
 MCAST_PORT = 1982
-MCAST_ADDR = MCAST_IP, MCAST_PORT
-
-
-class MCAST_MSG_TYPES:
-    SEARCH = 'M-SEARCH'
-    NOTIFY = 'NOTIFY'
+MCAST_ADDR = SimpleServiceDiscoveryProtocol.MULTICAST_ADDRESS, MCAST_PORT
 
 
 class MutableBoolean:
@@ -40,47 +32,35 @@ def set(self, value):
 
 @asyncio.coroutine
 def send_search_broadcast(transport, search_interval=30, _running=True):
+    request = SSDPRequest('M-SEARCH', headers=[
+        ('HOST', '%s:%s' % MCAST_ADDR),
+        ('MAN', '"ssdp:discover"'),
+        ('ST', 'wifi_bulb'),
+    ])
     while _running:
-        lines = ['%s * HTTP/1.1' % MCAST_MSG_TYPES.SEARCH]
-        lines += [
-            'HOST: %s:%s' % MCAST_ADDR,
-            'MAN: "ssdp:discover"',
-            'ST: wifi_bulb',
-        ]
-        msg = '\r\n'.join(lines)
-        logger.debug(">>> %s", msg)
-        transport.sendto(msg.encode(), MCAST_ADDR)
+        request.sendto(transport, MCAST_ADDR)
         yield from asyncio.sleep(search_interval)
 
 
-class YeelightProtocol(asyncio.DatagramProtocol):
+class YeelightProtocol(SimpleServiceDiscoveryProtocol):
     excluded_headers = ['DATE', 'EXT', 'SERVER', 'CACHE-CONTROL', 'LOCATION']
     location_patter = r'yeelight://(?P<ip>\d{1,3}(\.\d{1,3}){3}):(?P<port>\d+)'
 
     def __init__(self, bulbs, bulb_class=Bulb):
         self.bulbs = bulbs
         self.bulb_class = bulb_class
 
-    def datagram_received(self, data, addr):
-        msg = data.decode()
-        logger.debug("%s:%s> %s", addr + (msg,))
-
-        lines = msg.splitlines()
-        type, addr, status = lines[0].split()
-        if type == MCAST_MSG_TYPES.SEARCH:
-            return
-
-        data = '\n'.join(lines[1:])
-        headers = email.parser.Parser().parsestr(data)
-
+    @classmethod
+    def header_to_kwargs(cls, headers):
+        headers = dict(headers)
         location = headers.get('Location')
         cache_control = headers.get('Cache-Control', 'max-age=3600')
         headers = {
-            k: v for k, v in headers._headers
-            if k.upper() not in self.excluded_headers
+            k: v for k, v in headers.items()
+            if k.upper() not in cls.excluded_headers
         }
 
-        match = re.match(self.location_patter, location)
+        match = re.match(cls.location_patter, location)
         if match is None:
             raise YeelightError('Location does not match: %s' % location)
         ip = match.groupdict()['ip']
@@ -90,31 +70,31 @@ def datagram_received(self, data, addr):
 
         kwargs = dict(ip=ip, port=port, status_refresh_interv=max_age)
         kwargs.update(headers)
+        return kwargs
 
+    def request_received(self, request):
+        if request.method == 'M-SEARCH':
+            return
+        self.register_bulb(**self.header_to_kwargs(request.headers))
+
+    def response_received(self, response):
+        self.register_bulb(**self.header_to_kwargs(response.headers))
+
+    def register_bulb(self, **kwargs):
         idx = kwargs['id']
         if idx not in self.bulbs:
             self.bulbs[idx] = self.bulb_class(**kwargs)
         else:
             self.bulbs[idx].last_seen = time.time()
 
-    def error_received(self, exc):
-        if exc == errno.EAGAIN or exc == errno.EWOULDBLOCK:
-            logger.error('Error received: %s', exc)
-        else:
-            raise YeelightError("Unexpected connection error") from exc
-
-    def connection_lost(self, exc):
-        logger.error("Socket closed, stop the event loop")
-        loop = asyncio.get_event_loop()
-        loop.stop()
-
 
 @contextmanager
 def _unicast_socket():
     with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as ucast_socket:
         ucast_socket.bind(('', MCAST_PORT))
         fcntl.fcntl(ucast_socket, fcntl.F_SETFL, os.O_NONBLOCK)
-        group = socket.inet_aton(MCAST_IP)
+        group = socket.inet_aton(
+            SimpleServiceDiscoveryProtocol.MULTICAST_ADDRESS)
         mreq = struct.pack("4sl", group, socket.INADDR_ANY)
         ucast_socket.setsockopt(
             socket.IPPROTO_IP, socket.IP_ADD_MEMBERSHIP, mreq)