From fbd3c3d85a545f96f16954ff19ce0554655846cf Mon Sep 17 00:00:00 2001 From: William Pietri Date: Mon, 11 Dec 2017 16:14:52 -0800 Subject: [PATCH 1/2] Seeing if I can split the XMPP responsibilities from the vacuum responsibilities. --- sucks/__init__.py | 109 +++++++++++++++++++++---------------- tests/test_cli.py | 14 ----- tests/test_commands.py | 51 +++++++++++++++++ tests/test_ecovacs_xmpp.py | 15 +++++ tests/test_vacbot.py | 72 ++++-------------------- 5 files changed, 137 insertions(+), 124 deletions(-) create mode 100644 tests/test_commands.py create mode 100644 tests/test_ecovacs_xmpp.py diff --git a/sucks/__init__.py b/sucks/__init__.py index 8560364..59bb695 100644 --- a/sucks/__init__.py +++ b/sucks/__init__.py @@ -126,57 +126,37 @@ class EcoVacsAPI: return str(b64encode(result), 'utf8') -class VacBot(ClientXMPP): +class VacBot(): def __init__(self, user, domain, resource, secret, vacuum, continent): - ClientXMPP.__init__(self, user + '@' + domain, '0/' + resource + '/' + secret) - self.user = user - self.domain = domain - self.resource = resource self.vacuum = vacuum - self.continent = continent - self.credentials['authzid'] = user - self.add_event_handler("session_start", self.session_start) - - self.ready_flag = Event() self.clean_status = None self.charge_status = None self.battery_status = None - def wait_until_ready(self): - self.ready_flag.wait() + self.xmpp = EcoVacsXMPP(user, domain, resource, secret, continent) + self.xmpp.register_callback("CleanReport", self._handle_clean_report) + self.xmpp.register_callback("ChargeState", self._handle_charge_report) + self.xmpp.register_callback("BatteryInfo", self._handle_battery_report) + self.xmpp.register_callback("error", self._handle_error) - def session_start(self, event): - logging.debug("----------------- starting session ----------------") - logging.debug("event = {}".format(event)) - self.ready_flag.set() + def connect_and_wait_until_ready(self): + self.xmpp.connect_and_wait_until_ready() - self.__register_callback("CleanReport", self.handle_clean_report) - self.__register_callback("ChargeState", self.handle_charge_report) - self.__register_callback("BatteryInfo", self.handle_battery_report) - self.__register_callback("error", self.handle_error) + self.xmpp.schedule('Ping', 30, lambda: self.xmpp.send_ping(self._vacuum_adress()), repeat=True) - self.schedule('Ping', 30, self.send_ping, repeat=True) - - - def __register_callback(self, kind, function): - self.register_handler(Callback(kind, - MatchXPath('{jabber:client}iq/{com:ctl}query/{com:ctl}ctl[@td="' + kind + '"]'), - function)) - - - def handle_clean_report(self, iq): + def _handle_clean_report(self, iq): self.clean_status = iq.find('{com:ctl}query/{com:ctl}ctl/{com:ctl}clean').get('type') logging.debug("*** clean_status = " + self.clean_status) - def handle_battery_report(self, iq): + def _handle_battery_report(self, iq): try: self.battery_status = float(iq.find('{com:ctl}query/{com:ctl}ctl/{com:ctl}battery').get('power')) / 100 logging.debug("*** battery_status = {:.0%}".format(self.battery_status)) except ValueError: logging.warning("couldn't parse battery status " + ET.tostring(iq)) - def handle_charge_report(self, iq): + def _handle_charge_report(self, iq): report = iq.find('{com:ctl}query/{com:ctl}ctl/{com:ctl}charge').get('type') if report.lower() == 'going': self.charge_status = 'returning' @@ -188,45 +168,78 @@ class VacBot(ClientXMPP): logging.warning("Unknown charging status '" + report + "'") logging.debug("*** charge_status = " + self.charge_status) - def handle_error(self, iq): + def _handle_error(self, iq): error = iq.find('{com:ctl}query/{com:ctl}ctl').get('error') error_no = iq.find('{com:ctl}query/{com:ctl}ctl').get('errno') logging.debug("*** error = " + error_no + " " + error) + def _vacuum_adress(self): + return self.vacuum['did'] + '@' + self.vacuum['class'] + '.ecorobot.net/atom' + def send_command(self, xml): - c = self.wrap_command(xml) + self.xmpp.send_command(xml, self._vacuum_adress()) + + def run(self, action): + self.send_command(action.to_xml()) + action.wait_for_completion(self) + + def disconnect(self, wait=False): + self.xmpp.disconnect(wait=wait) + + +class EcoVacsXMPP(ClientXMPP): + def __init__(self, user, domain, resource, secret, continent): + ClientXMPP.__init__(self, user + '@' + domain, '0/' + resource + '/' + secret) + + self.user = user + self.domain = domain + self.resource = resource + self.continent = continent + self.credentials['authzid'] = user + self.add_event_handler("session_start", self.session_start) + + self.ready_flag = Event() + + def wait_until_ready(self): + self.ready_flag.wait() + + def session_start(self, event): + logging.debug("----------------- starting session ----------------") + logging.debug("event = {}".format(event)) + self.ready_flag.set() + + def register_callback(self, kind, function): + self.register_handler(Callback(kind, + MatchXPath('{jabber:client}iq/{com:ctl}query/{com:ctl}ctl[@td="' + kind + '"]'), + function)) + + def send_command(self, xml, recipient): + c = self._wrap_command(xml, recipient) logging.debug('Sending command {0}'.format(c)) c.send() - def wrap_command(self, ctl): - q = self.make_iq_query(xmlns=u'com:ctl', ito=self.__vacuum_adress(), ifrom=self.__my_address()) + def _wrap_command(self, ctl, recipient): + q = self.make_iq_query(xmlns=u'com:ctl', ito=recipient, ifrom=self._my_address()) q['type'] = 'set' for child in q.xml: if child.tag.endswith('query'): child.append(ctl) return q - def send_ping(self): - q = self.make_iq_get(ito=self.__vacuum_adress(), ifrom=self.__my_address()) + def _my_address(self): + return self.user + '@' + self.domain + '/' + self.resource + + def send_ping(self, to): + q = self.make_iq_get(ito=to, ifrom=self._my_address()) q.xml.append(ET.Element('ping', {'xmlns': 'urn:xmpp:ping'})) logging.debug("*** sending ping ***") q.send() - def __my_address(self): - return self.user + '@' + self.domain + '/' + self.resource - - def __vacuum_adress(self): - return self.vacuum['did'] + '@' + self.vacuum['class'] + '.ecorobot.net/atom' - def connect_and_wait_until_ready(self): self.connect(('msg-{}.ecouser.net'.format(self.continent), '5223')) self.process() self.wait_until_ready() - def run(self, action): - self.send_command(action.to_xml()) - action.wait_for_completion(self) - class VacBotCommand: def __init__(self, name, args=None, wait=None, terminal=False): diff --git a/tests/test_cli.py b/tests/test_cli.py index 8d55c3c..035e1e3 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -61,20 +61,6 @@ def test_continent_for_country(): assert_equal(continent_for_country('fr'), 'eu') -def test_wrap_command(): - v = VacBot('20170101abcdefabcdefa', 'ecouser.net', 'abcdef12', 'A1b2C3d4efghijklmNOPQrstuvwxyz12', - {"did": "E0000000001234567890", "class": "126", "nick": "bob"}, 'na') - c = str(v.wrap_command(Clean(1).to_xml())) - assert_true(re.search(r'from="20170101abcdefabcdefa@ecouser.net/abcdef12"', c)) - assert_true(re.search(r'to="E0000000001234567890@126.ecorobot.net/atom"', c)) - - -def test_model_variation(): - v = VacBot('20170101abcdefabcdefa', 'ecouser.net', 'abcdef12', 'A1b2C3d4efghijklmNOPQrstuvwxyz12', - {"did": "E0000000001234567890", "class": "141", "nick": "bob"}, 'na') - c = str(v.wrap_command(Clean(1).to_xml())) - assert_true(re.search(r'to="E0000000001234567890@141.ecorobot.net/atom"', c)) - def test_main_api_setup(): with requests_mock.mock() as m: diff --git a/tests/test_commands.py b/tests/test_commands.py new file mode 100644 index 0000000..e8d9c35 --- /dev/null +++ b/tests/test_commands.py @@ -0,0 +1,51 @@ +from xml.etree import ElementTree + +from nose.tools import * + +from sucks import * + + +def test_custom_command(): + # Ensure a custom-built command generates the expected XML payload + c = VacBotCommand('CustomCommand', {'type': 'customtype'}) + assert_equals(ElementTree.tostring(c.to_xml()), + + b'') + + +def test_custom_command_noargs(): + # Ensure a custom-built command with no args generates XML without an args element + c = VacBotCommand('CustomCommand') + assert_equals(ElementTree.tostring(c.to_xml()), + b'') + + +def test_clean_command(): + c = Clean(10) + assert_equals(c.terminal, False) + assert_equals(c.wait, 10) + assert_equals(ElementTree.tostring(c.to_xml()), + b'') # protocol has attribs in other order + + +def test_edge_command(): + # called Edge because that's what the UI uses, even though the protocol is different + c = Edge(10) + assert_equals(c.terminal, False) + assert_equals(c.wait, 10) + assert_equals(ElementTree.tostring(c.to_xml()), + b'') # protocol has attribs in other order + + +def test_charge_command(): + c = Charge() + assert_equals(c.terminal, True) + assert_equals(ElementTree.tostring(c.to_xml()), + b'') + + +def test_stop_command(): + c = Stop() + assert_equals(c.terminal, True) + assert_equals(ElementTree.tostring(c.to_xml()), + b'') diff --git a/tests/test_ecovacs_xmpp.py b/tests/test_ecovacs_xmpp.py new file mode 100644 index 0000000..d0576f2 --- /dev/null +++ b/tests/test_ecovacs_xmpp.py @@ -0,0 +1,15 @@ +from re import search + +from nose.tools import * + +from sucks import * + + +# There are few tests for the XMPP stuff here because it's relatively complicated to test given +# the library's design and its multithreaded nature and lack of explicit testing support. + +def test_wrap_command(): + x = EcoVacsXMPP('20170101abcdefabcdefa', 'ecouser.net', 'abcdef12', 'A1b2C3d4efghijklmNOPQrstuvwxyz12', 'na') + c = str(x._wrap_command(Clean(1).to_xml(), 'E0000000001234567890@126.ecorobot.net/atom')) + assert_true(search(r'from="20170101abcdefabcdefa@ecouser.net/abcdef12"', c)) + assert_true(search(r'to="E0000000001234567890@126.ecorobot.net/atom"', c)) diff --git a/tests/test_vacbot.py b/tests/test_vacbot.py index 95e6555..a918fdb 100644 --- a/tests/test_vacbot.py +++ b/tests/test_vacbot.py @@ -1,70 +1,18 @@ -from xml.etree import ElementTree - -from re import search from nose.tools import * from sucks import * -# There are few tests for the XMPP stuff here because a) it's relatively complicated to test given -# the library's design and its multithreaded nature, and b) I'm manually testing every change anyhow, -# as it's not clear how the robot really behaves. - -def test_custom_command(): - # Ensure a custom-built command generates the expected XML payload - c = VacBotCommand('CustomCommand', {'type': 'customtype'}) - assert_equals(ElementTree.tostring(c.to_xml()), - - b'') - -def test_custom_command_noargs(): - # Ensure a custom-built command with no args generates XML without an args element - c = VacBotCommand('CustomCommand') - assert_equals(ElementTree.tostring(c.to_xml()), - b'') - - -def test_clean_command(): - c = Clean(10) - assert_equals(c.terminal, False) - assert_equals(c.wait, 10) - assert_equals(ElementTree.tostring(c.to_xml()), - b'') # protocol has attribs in other order - - -def test_edge_command(): - # called Edge because that's what the UI uses, even though the protocol is different - c = Edge(10) - assert_equals(c.terminal, False) - assert_equals(c.wait, 10) - assert_equals(ElementTree.tostring(c.to_xml()), - b'') # protocol has attribs in other order - - -def test_charge_command(): - c = Charge() - assert_equals(c.terminal, True) - assert_equals(ElementTree.tostring(c.to_xml()), - b'') - - -def test_stop_command(): - c = Stop() - assert_equals(c.terminal, True) - assert_equals(ElementTree.tostring(c.to_xml()), - b'') - - -def test_wrap_command(): - v = VacBot('20170101abcdefabcdefa', 'ecouser.net', 'abcdef12', 'A1b2C3d4efghijklmNOPQrstuvwxyz12', - {"did": "E0000000001234567890", "class": "126", "nick": "bob"}, 'na') - c = str(v.wrap_command(Clean(1).to_xml())) - assert_true(search(r'from="20170101abcdefabcdefa@ecouser.net/abcdef12"', c)) - assert_true(search(r'to="E0000000001234567890@126.ecorobot.net/atom"', c)) +def test_bot_address(): + v = vacbot_for_bot({"did": "E0000000001234567890", "class": "126", "nick": "bob"}) + assert_equals('E0000000001234567890@126.ecorobot.net/atom', v._vacuum_adress()) def test_model_variation(): - v = VacBot('20170101abcdefabcdefa', 'ecouser.net', 'abcdef12', 'A1b2C3d4efghijklmNOPQrstuvwxyz12', - {"did": "E0000000001234567890", "class": "141", "nick": "bob"}, 'na') - c = str(v.wrap_command(Clean(1).to_xml())) - assert_true(search(r'to="E0000000001234567890@141.ecorobot.net/atom"', c)) + v = vacbot_for_bot({"did": "E0000000001234567890", "class": "141", "nick": "bob"}) + assert_equals('E0000000001234567890@141.ecorobot.net/atom', v._vacuum_adress()) + + +def vacbot_for_bot(bot): + return VacBot('20170101abcdefabcdefa', 'ecouser.net', 'abcdef12', 'A1b2C3d4efghijklmNOPQrstuvwxyz12', + bot, 'na') From d6e4f7d5c50b194909cf4bdb2b7b4bec33d83b04 Mon Sep 17 00:00:00 2001 From: William Pietri Date: Tue, 12 Dec 2017 11:49:10 -0800 Subject: [PATCH 2/2] Events now flow to Vacbot as reasonably sane dicts. --- setup.py | 3 +- sucks/__init__.py | 60 +++++++++++++++++++++++++--------- tests/test_ecovacs_xmpp.py | 47 ++++++++++++++++++++++++++- tests/test_vacbot.py | 66 +++++++++++++++++++++++++++++++++++--- 4 files changed, 154 insertions(+), 22 deletions(-) diff --git a/setup.py b/setup.py index 443e77d..8e85404 100644 --- a/setup.py +++ b/setup.py @@ -67,7 +67,8 @@ setup( 'click>=6', 'requests>=2.18', 'pycryptodome>=3.4', - 'pycountry-convert>=0.5' + 'pycountry-convert>=0.5', + 'stringcase>=1.2' ], # List additional groups of dependencies here (e.g. development diff --git a/sucks/__init__.py b/sucks/__init__.py index 59bb695..c009139 100644 --- a/sucks/__init__.py +++ b/sucks/__init__.py @@ -7,6 +7,7 @@ from threading import Event import click import requests +import stringcase from sleekxmpp import ClientXMPP, Callback, MatchXPath from sleekxmpp.xmlstream import ET @@ -135,34 +136,38 @@ class VacBot(): self.battery_status = None self.xmpp = EcoVacsXMPP(user, domain, resource, secret, continent) - self.xmpp.register_callback("CleanReport", self._handle_clean_report) - self.xmpp.register_callback("ChargeState", self._handle_charge_report) - self.xmpp.register_callback("BatteryInfo", self._handle_battery_report) self.xmpp.register_callback("error", self._handle_error) + self.xmpp.subscribe_to_ctls(self._handle_ctl) def connect_and_wait_until_ready(self): self.xmpp.connect_and_wait_until_ready() - self.xmpp.schedule('Ping', 30, lambda: self.xmpp.send_ping(self._vacuum_adress()), repeat=True) + self.xmpp.schedule('Ping', 30, lambda: self.xmpp.send_ping(self._vacuum_address()), repeat=True) - def _handle_clean_report(self, iq): - self.clean_status = iq.find('{com:ctl}query/{com:ctl}ctl/{com:ctl}clean').get('type') + def _handle_ctl(self, ctl): + method = '_handle_' + ctl['event'] + if hasattr(self, method): + getattr(self, method)(ctl) + + + def _handle_clean_report(self, event): + self.clean_status = event['type'] logging.debug("*** clean_status = " + self.clean_status) - def _handle_battery_report(self, iq): + def _handle_battery_info(self, iq): try: - self.battery_status = float(iq.find('{com:ctl}query/{com:ctl}ctl/{com:ctl}battery').get('power')) / 100 + self.battery_status = float(iq['power']) / 100 logging.debug("*** battery_status = {:.0%}".format(self.battery_status)) except ValueError: logging.warning("couldn't parse battery status " + ET.tostring(iq)) - def _handle_charge_report(self, iq): - report = iq.find('{com:ctl}query/{com:ctl}ctl/{com:ctl}charge').get('type') - if report.lower() == 'going': + def _handle_charge_state(self, event): + report = event['type'] + if report == 'going': self.charge_status = 'returning' - elif report.lower() == 'slotcharging': + elif report == 'slot_charging': self.charge_status = 'charging' - elif report.lower() == 'idle': + elif report == 'idle': self.charge_status = 'idle' else: logging.warning("Unknown charging status '" + report + "'") @@ -173,11 +178,11 @@ class VacBot(): error_no = iq.find('{com:ctl}query/{com:ctl}ctl').get('errno') logging.debug("*** error = " + error_no + " " + error) - def _vacuum_adress(self): + def _vacuum_address(self): return self.vacuum['did'] + '@' + self.vacuum['class'] + '.ecorobot.net/atom' def send_command(self, xml): - self.xmpp.send_command(xml, self._vacuum_adress()) + self.xmpp.send_command(xml, self._vacuum_address()) def run(self, action): self.send_command(action.to_xml()) @@ -198,6 +203,7 @@ class EcoVacsXMPP(ClientXMPP): self.credentials['authzid'] = user self.add_event_handler("session_start", self.session_start) + self.ctl_subscribers = [] self.ready_flag = Event() def wait_until_ready(self): @@ -206,8 +212,32 @@ class EcoVacsXMPP(ClientXMPP): def session_start(self, event): logging.debug("----------------- starting session ----------------") logging.debug("event = {}".format(event)) + self.register_handler(Callback("general", + MatchXPath('{jabber:client}iq/{com:ctl}query/{com:ctl}'), + self._handle_ctl)) + self.ready_flag.set() + def subscribe_to_ctls(self, function): + self.ctl_subscribers.append(function) + + + def _handle_ctl(self, message): + the_good_part = message.get_payload()[0][0] + as_dict = self._ctl_to_dict(the_good_part) + for s in self.ctl_subscribers: + s(as_dict) + + def _ctl_to_dict(self, xml): + result = xml.attrib.copy() + result['event'] = result.pop('td') + if xml: + result.update(xml[0].attrib) + + for key in result: + result[key] = stringcase.snakecase(result[key]) + return result + def register_callback(self, kind, function): self.register_handler(Callback(kind, MatchXPath('{jabber:client}iq/{com:ctl}query/{com:ctl}ctl[@td="' + kind + '"]'), diff --git a/tests/test_ecovacs_xmpp.py b/tests/test_ecovacs_xmpp.py index d0576f2..8bcc0ad 100644 --- a/tests/test_ecovacs_xmpp.py +++ b/tests/test_ecovacs_xmpp.py @@ -9,7 +9,52 @@ from sucks import * # the library's design and its multithreaded nature and lack of explicit testing support. def test_wrap_command(): - x = EcoVacsXMPP('20170101abcdefabcdefa', 'ecouser.net', 'abcdef12', 'A1b2C3d4efghijklmNOPQrstuvwxyz12', 'na') + x = make_ecovacs_xmpp() c = str(x._wrap_command(Clean(1).to_xml(), 'E0000000001234567890@126.ecorobot.net/atom')) assert_true(search(r'from="20170101abcdefabcdefa@ecouser.net/abcdef12"', c)) assert_true(search(r'to="E0000000001234567890@126.ecorobot.net/atom"', c)) + + +def test_subscribe_to_ctls(): + response = None + + def save_response(value): + nonlocal response + response = value + + x = make_ecovacs_xmpp() + + query = x.make_iq_query() + query.set_payload( + ET.fromstring(' ')) + + x.subscribe_to_ctls(save_response) + x._handle_ctl(query) + assert_dict_equal(response, {'event': 'clean_report', 'type': 'auto'}) + + +def test_xml_to_dict(): + x = make_ecovacs_xmpp() + + assert_dict_equal( + x._ctl_to_dict(make_ctl(' ')), + {'event': 'clean_report', 'type': 'auto'}) + assert_dict_equal( + x._ctl_to_dict(make_ctl(' ')), + {'event': 'clean_report', 'type': 'auto', 'speed': 'strong'}) + + assert_dict_equal( + x._ctl_to_dict(make_ctl('')), + {'event': 'battery_info', 'power': '095'}) + + assert_dict_equal( + x._ctl_to_dict(make_ctl('# ')), + {'event': 'life_span', 'type': 'brush', 'val': '099', 'total': '365'}) + + +def make_ecovacs_xmpp(): + return EcoVacsXMPP('20170101abcdefabcdefa', 'ecouser.net', 'abcdef12', 'A1b2C3d4efghijklmNOPQrstuvwxyz12', 'na') + + +def make_ctl(string): + return ET.fromstring('' + string + '')[0] diff --git a/tests/test_vacbot.py b/tests/test_vacbot.py index a918fdb..8da9ab3 100644 --- a/tests/test_vacbot.py +++ b/tests/test_vacbot.py @@ -3,16 +3,72 @@ from nose.tools import * from sucks import * + +def test_handle_clean_report(): + v = a_vacbot() + assert_equals(None, v.clean_status) + + v._handle_ctl({'event': 'clean_report', 'type': 'auto', 'speed': 'strong'}) + assert_equals('auto', v.clean_status) + + +def test_handle_charge_state(): + v = a_vacbot() + assert_equals(None, v.clean_status) + + v._handle_ctl({'event': 'charge_state', 'type': 'going'}) + assert_equals('returning', v.charge_status) + + v._handle_ctl({'event': 'charge_state', 'type': 'slot_charging'}) + assert_equals('charging', v.charge_status) + + v._handle_ctl({'event': 'charge_state', 'type': 'idle'}) + assert_equals('idle', v.charge_status) + + +def test_handle_battery_info(): + v = a_vacbot() + assert_equals(None, v.battery_status) + + v._handle_ctl({'event': 'battery_info', 'power': '100'}) + assert_equals(1.0, v.battery_status) + + v._handle_ctl({'event': 'battery_info', 'power': '095'}) + assert_equals(0.95, v.battery_status) + + v._handle_ctl({'event': 'battery_info', 'power': '000'}) + assert_equals(0.0, v.battery_status) + + +def test_handle_unknown_ctl(): + v = a_vacbot() + v._handle_ctl({'event': 'weird_and_unknown_event', 'type': 'pretty_weird'}) + # as long as it doesn't blow up, that's fine + + +# as-yet unhandled messages: +# +# +# +# +# +# +# +# plus errors! + def test_bot_address(): - v = vacbot_for_bot({"did": "E0000000001234567890", "class": "126", "nick": "bob"}) - assert_equals('E0000000001234567890@126.ecorobot.net/atom', v._vacuum_adress()) + v = a_vacbot(bot={"did": "E0000000001234567890", "class": "126", "nick": "bob"}) + assert_equals('E0000000001234567890@126.ecorobot.net/atom', v._vacuum_address()) def test_model_variation(): - v = vacbot_for_bot({"did": "E0000000001234567890", "class": "141", "nick": "bob"}) - assert_equals('E0000000001234567890@141.ecorobot.net/atom', v._vacuum_adress()) + v = a_vacbot(bot={"did": "E0000000001234567890", "class": "141", "nick": "bob"}) + assert_equals('E0000000001234567890@141.ecorobot.net/atom', v._vacuum_address()) -def vacbot_for_bot(bot): + +def a_vacbot(bot=None): + if bot is None: + bot = {"did": "E0000000001234567890", "class": "126", "nick": "bob"} return VacBot('20170101abcdefabcdefa', 'ecouser.net', 'abcdef12', 'A1b2C3d4efghijklmNOPQrstuvwxyz12', bot, 'na')