Skip to content

Commit 74e0598

Browse files
codingjoecodingjoe
committed
Isolate SSDP into separate module
1 parent b2f0a3a commit 74e0598

File tree

7 files changed

+302
-104
lines changed

7 files changed

+302
-104
lines changed

‎setup.cfg‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,5 +23,9 @@ keywords =
2323
packages =
2424
yeelib
2525

26+
[pycodestyle]
27+
max-line-length=99
28+
exclude = fixtures.py
29+
2630
[pydocstyle]
27-
add_ignore = D1
31+
add_ignore = D1,D401

‎tests/fixtures.py‎

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
1+
request = b"""NOTIFY * HTTP/1.1
2+
Host: 239.255.255.250:1982
3+
Cache-Control: max-age=3600
4+
Location: yeelight://192.168.1.239:55443
5+
NTS: ssdp:alive
6+
Server: POSIX, UPnP/1.0 YGLC/1
7+
id: 0x000000000015243f
8+
model: color
9+
fw_ver: 18
10+
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
11+
power: on
12+
bright: 100
13+
color_mode: 2
14+
ct: 4000
15+
rgb: 16711680
16+
hue: 100
17+
sat: 35
18+
name: my_bulb""".replace(b'\n', b'\r\n')
19+
20+
response = b"""HTTP/1.1 200 OK
21+
Cache-Control: max-age=3600
22+
Date:
23+
Ext:
24+
Location: yeelight://192.168.1.239:55443
25+
Server: POSIX UPnP/1.0 YGLC/1
26+
id: 0x000000000015243f
27+
model: color
28+
fw_ver: 18
29+
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
30+
power: on
31+
bright: 100
32+
color_mode: 2
33+
ct: 4000
34+
rgb: 16711680
35+
hue: 100
36+
sat: 35
37+
name: my_bulb""".replace(b'\n', b'\r\n')
38+
39+
response_wrong_location = b"""HTTP/1.1 200 OK
40+
Cache-Control: max-age=3600
41+
Date:
42+
Ext:
43+
Location: yeelight://not.an.ip:55443
44+
Server: POSIX UPnP/1.0 YGLC/1""".replace(b'\n', b'\r\n')

‎tests/test_bulbs.py‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -26,8 +26,9 @@ def test_kwargs(self):
2626
'id': '0x000000000015243f',
2727
'model': 'color',
2828
'fw_ver': 18,
29-
'support': 'get_prop set_default set_power toggle set_bright start_cf stop_cf set_scene'
30-
'cron_add cron_get cron_del set_ct_abx set_rgb',
29+
'support': 'get_prop set_default set_power toggle set_bright'
30+
' start_cf stop_cf set_scene cron_add cron_get'
31+
' cron_del set_ct_abx set_rgb',
3132
'power': 'on',
3233
'bright': 100,
3334
'color_mode': 2,
@@ -41,5 +42,5 @@ def test_kwargs(self):
4142
with Bulb(*self.bulb_addr, **kwargs) as b:
4243
assert b.fw_ver == 18
4344
assert b.support == ['get_prop', 'set_default', 'set_power', 'toggle', 'set_bright',
44-
'start_cf', 'stop_cf', 'set_scenecron_add', 'cron_get',
45+
'start_cf', 'stop_cf', 'set_scene', 'cron_add', 'cron_get',
4546
'cron_del', 'set_ct_abx', 'set_rgb']

‎tests/test_discover.py‎

Lines changed: 7 additions & 52 deletions
Original file line numberDiff line numberDiff line change
@@ -5,83 +5,38 @@
55
from yeelib.discover import YeelightProtocol, search_bulbs
66
from yeelib.exceptions import YeelightError
77

8-
9-
notify = b"""NOTIFY * HTTP/1.1
10-
Host: 239.255.255.250:1982
11-
Cache-Control: max-age=3600
12-
Location: yeelight://192.168.1.239:55443
13-
NTS: ssdp:alive
14-
Server: POSIX, UPnP/1.0 YGLC/1
15-
id: 0x000000000015243f
16-
model: color
17-
fw_ver: 18
18-
support: get_prop set_default set_power toggle set_bright start_cf stop_cf set_scene
19-
cron_add cron_get cron_del set_ct_abx set_rgb
20-
power: on
21-
bright: 100
22-
color_mode: 2
23-
ct: 4000
24-
rgb: 16711680
25-
hue: 100
26-
sat: 35
27-
name: my_bulb"""
28-
29-
mcast = b"""HTTP/1.1 200 OK
30-
Cache-Control: max-age=3600
31-
Date:
32-
Ext:
33-
Location: yeelight://192.168.1.239:55443
34-
Server: POSIX UPnP/1.0 YGLC/1
35-
id: 0x000000000015243f
36-
model: color
37-
fw_ver: 18
38-
support: get_prop set_default set_power toggle set_bright start_cf stop_cf set_scene
39-
cron_add cron_get cron_del set_ct_abx set_rgb
40-
power: on
41-
bright: 100
42-
color_mode: 2
43-
ct: 4000
44-
rgb: 16711680
45-
hue: 100
46-
sat: 35
47-
name: my_bulb"""
48-
49-
wrong_location = b"""HTTP/1.1 200 OK
50-
Cache-Control: max-age=3600
51-
Date:
52-
Ext:
53-
Location: yeelight://not.an.ip:55443
54-
Server: POSIX UPnP/1.0 YGLC/1"""
8+
from . import fixtures
559

5610

5711
class TestYeelightProtocoll:
5812
def test_notify(self, ):
5913
bulbs = {}
6014
p = YeelightProtocol(bulbs=bulbs)
61-
p.datagram_received(data=notify, addr=('192.168.1.239', 1982))
15+
p.datagram_received(data=fixtures.request, addr=('192.168.1.239', 1982))
6216
assert len(bulbs) == 1
6317
assert bulbs['0x000000000015243f'].ip == '192.168.1.239'
6418

6519
def test_mcast(self, ):
6620
bulbs = {}
6721
p = YeelightProtocol(bulbs=bulbs)
68-
p.datagram_received(data=mcast, addr=('192.168.1.239', 1982))
22+
p.datagram_received(data=fixtures.response, addr=('192.168.1.239', 1982))
6923
assert len(bulbs) == 1
7024
assert bulbs['0x000000000015243f'].ip == '192.168.1.239'
7125

7226
def test_duplicate(self):
7327
bulbs = {}
7428
p = YeelightProtocol(bulbs=bulbs)
75-
p.datagram_received(data=notify, addr=('192.168.1.239', 1982))
76-
p.datagram_received(data=notify, addr=('192.168.1.239', 1982))
29+
p.datagram_received(data=fixtures.request, addr=('192.168.1.239', 1982))
30+
p.datagram_received(data=fixtures.request, addr=('192.168.1.239', 1982))
7731
assert len(bulbs) == 1
7832
assert bulbs['0x000000000015243f'].ip == '192.168.1.239'
7933

8034
def test_wrong_location(self):
8135
bulbs = {}
8236
p = YeelightProtocol(bulbs=bulbs)
8337
with pytest.raises(YeelightError) as e:
84-
p.datagram_received(data=wrong_location, addr=('192.168.1.239', 1982))
38+
p.datagram_received(data=fixtures.response_wrong_location,
39+
addr=('192.168.1.239', 1982))
8540
assert 'Location does not match: yeelight://not.an.ip:55443' in str(e)
8641

8742

‎tests/test_upnp.py‎

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
1+
import pytest
2+
3+
from yeelib.upnp import SSDPMessage, SSDPResponse, SSDPRequest
4+
from . import fixtures
5+
6+
7+
class TestSSDPMessage:
8+
def test_headers_copy(self):
9+
headers = [('Cache-Control', 'max-age=3600')]
10+
msg = SSDPMessage(headers=headers)
11+
assert msg.headers == headers
12+
assert msg.headers is not headers
13+
14+
def test_headers_dict(self):
15+
headers = {'Cache-Control': 'max-age=3600'}
16+
msg = SSDPMessage(headers=headers)
17+
assert msg.headers == [('Cache-Control', 'max-age=3600')]
18+
19+
def test_headers_none(self):
20+
msg = SSDPMessage(headers=None)
21+
assert msg.headers == []
22+
23+
def test_parse(self):
24+
with pytest.raises(NotImplementedError):
25+
SSDPMessage.parse('')
26+
27+
def test_parse_headers(self):
28+
headers = SSDPMessage.parse_headers('Cache-Control: max-age=3600')
29+
assert headers == [('Cache-Control', 'max-age=3600')]
30+
31+
32+
class TestSSDPResponse:
33+
def test_parse(self):
34+
response = SSDPResponse.parse(fixtures.response.decode())
35+
assert response.status_code == 200
36+
assert response.reason == 'OK'
37+
38+
39+
class TestSSDPRequest:
40+
def test_parse(self):
41+
request = SSDPRequest.parse(fixtures.request.decode())
42+
assert request.method == 'NOTIFY'
43+
assert request.uri == '*'
44+
45+
def test_str(self):
46+
request = SSDPRequest('NOTIFY', '*', headers=[('Cache-Control', 'max-age=3600')])
47+
assert str(request) == (
48+
'NOTIFY * HTTP/1.1\n'
49+
'Cache-Control: max-age=3600'
50+
)
51+
52+
def test_bytes(self):
53+
request = SSDPRequest.parse(fixtures.request.decode())
54+
assert bytes(request) == fixtures.request

‎yeelib/discover.py‎

Lines changed: 28 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,4 @@
11
import asyncio
2-
import email.parser
3-
import errno
42
import fcntl
53
import logging
64
import os
@@ -11,21 +9,15 @@
119
from contextlib import contextmanager
1210

1311
from yeelib.bulbs import Bulb
14-
12+
from yeelib.upnp import SSDPRequest, SimpleServiceDiscoveryProtocol
1513
from .exceptions import YeelightError
1614

17-
__all__ = ('search_bulbs',)
15+
__all__ = ('search_bulbs', 'YeelightProtocol')
1816

1917
logger = logging.getLogger('yeelib')
2018

21-
MCAST_IP = '239.255.255.250'
2219
MCAST_PORT = 1982
23-
MCAST_ADDR = MCAST_IP, MCAST_PORT
24-
25-
26-
class MCAST_MSG_TYPES:
27-
SEARCH = 'M-SEARCH'
28-
NOTIFY = 'NOTIFY'
20+
MCAST_ADDR = SimpleServiceDiscoveryProtocol.MULTICAST_ADDRESS, MCAST_PORT
2921

3022

3123
class MutableBoolean:
@@ -40,47 +32,35 @@ def set(self, value):
4032

4133
@asyncio.coroutine
4234
def send_search_broadcast(transport, search_interval=30, _running=True):
35+
request = SSDPRequest('M-SEARCH', headers=[
36+
('HOST', '%s:%s' % MCAST_ADDR),
37+
('MAN', '"ssdp:discover"'),
38+
('ST', 'wifi_bulb'),
39+
])
4340
while _running:
44-
lines = ['%s * HTTP/1.1' % MCAST_MSG_TYPES.SEARCH]
45-
lines += [
46-
'HOST: %s:%s' % MCAST_ADDR,
47-
'MAN: "ssdp:discover"',
48-
'ST: wifi_bulb',
49-
]
50-
msg = '\r\n'.join(lines)
51-
logger.debug(">>> %s", msg)
52-
transport.sendto(msg.encode(), MCAST_ADDR)
41+
request.sendto(transport, MCAST_ADDR)
5342
yield from asyncio.sleep(search_interval)
5443

5544

56-
class YeelightProtocol(asyncio.DatagramProtocol):
45+
class YeelightProtocol(SimpleServiceDiscoveryProtocol):
5746
excluded_headers = ['DATE', 'EXT', 'SERVER', 'CACHE-CONTROL', 'LOCATION']
5847
location_patter = r'yeelight://(?P<ip>\d{1,3}(\.\d{1,3}){3}):(?P<port>\d+)'
5948

6049
def __init__(self, bulbs, bulb_class=Bulb):
6150
self.bulbs = bulbs
6251
self.bulb_class = bulb_class
6352

64-
def datagram_received(self, data, addr):
65-
msg = data.decode()
66-
logger.debug("%s:%s> %s", addr + (msg,))
67-
68-
lines = msg.splitlines()
69-
type, addr, status = lines[0].split()
70-
if type == MCAST_MSG_TYPES.SEARCH:
71-
return
72-
73-
data = '\n'.join(lines[1:])
74-
headers = email.parser.Parser().parsestr(data)
75-
53+
@classmethod
54+
def header_to_kwargs(cls, headers):
55+
headers = dict(headers)
7656
location = headers.get('Location')
7757
cache_control = headers.get('Cache-Control', 'max-age=3600')
7858
headers = {
79-
k: v for k, v in headers._headers
80-
if k.upper() not in self.excluded_headers
59+
k: v for k, v in headers.items()
60+
if k.upper() not in cls.excluded_headers
8161
}
8262

83-
match = re.match(self.location_patter, location)
63+
match = re.match(cls.location_patter, location)
8464
if match is None:
8565
raise YeelightError('Location does not match: %s' % location)
8666
ip = match.groupdict()['ip']
@@ -90,31 +70,31 @@ def datagram_received(self, data, addr):
9070

9171
kwargs = dict(ip=ip, port=port, status_refresh_interv=max_age)
9272
kwargs.update(headers)
73+
return kwargs
9374

75+
def request_received(self, request):
76+
if request.method == 'M-SEARCH':
77+
return
78+
self.register_bulb(**self.header_to_kwargs(request.headers))
79+
80+
def response_received(self, response):
81+
self.register_bulb(**self.header_to_kwargs(response.headers))
82+
83+
def register_bulb(self, **kwargs):
9484
idx = kwargs['id']
9585
if idx not in self.bulbs:
9686
self.bulbs[idx] = self.bulb_class(**kwargs)
9787
else:
9888
self.bulbs[idx].last_seen = time.time()
9989

100-
def error_received(self, exc):
101-
if exc == errno.EAGAIN or exc == errno.EWOULDBLOCK:
102-
logger.error('Error received: %s', exc)
103-
else:
104-
raise YeelightError("Unexpected connection error") from exc
105-
106-
def connection_lost(self, exc):
107-
logger.error("Socket closed, stop the event loop")
108-
loop = asyncio.get_event_loop()
109-
loop.stop()
110-
11190

11291
@contextmanager
11392
def _unicast_socket():
11493
with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as ucast_socket:
11594
ucast_socket.bind(('', MCAST_PORT))
11695
fcntl.fcntl(ucast_socket, fcntl.F_SETFL, os.O_NONBLOCK)
117-
group = socket.inet_aton(MCAST_IP)
96+
group = socket.inet_aton(
97+
SimpleServiceDiscoveryProtocol.MULTICAST_ADDRESS)
11898
mreq = struct.pack("4sl", group, socket.INADDR_ANY)
11999
ucast_socket.setsockopt(
120100
socket.IPPROTO_IP, socket.IP_ADD_MEMBERSHIP, mreq)

0 commit comments

Comments
 (0)