VacBotCommand class updated and more commands added

sucks/__init__.py

@@ -272,7 +272,38 @@ class EcoVacsXMPP(ClientXMPP):
 
 
 class VacBotCommand:
-    def __init__(self, name, args=None, wait=None, terminal=False):
+
+    CLEAN_MODE ={
+        'auto': 'auto',
+        'edge': 'border',
+        'spot': 'spot',
+        'single_room': 'singleroom',
+        'stop': 'stop'
+    }
+    FAN_SPEED = {
+        'normal': 'standard',
+        'high': 'strong'
+    }
+    CHARGE_MODE = {
+        'return': 'go',
+        'returning': 'Going',
+        'charging': 'SlotCharging',
+        'idle': 'Idle'
+    }
+    COMPONENT = {
+        'main_brush': 'Brush',
+        'side_brush': 'SideBrush',
+        'filter': 'DustCaseHeap'
+    }
+    ACTION = {
+        'forward': 'forward',
+        'left': 'SpinLeft',
+        'right': 'SpinRight',
+        'turn_around': 'TurnAround',
+        'stop': 'stop'
+    }
+
+    def __init__(self, name, args={}, wait=None, terminal=False):
         self.name = name
         self.args = args
         self.wait = wait
@@ -285,9 +316,12 @@ class VacBotCommand:
 
     def to_xml(self):
         ctl = ET.Element('ctl', {'td': self.name})
-        if self.args:
+        for key, value in self.args.items():
-            inner = ET.Element(self.name.lower(), self.args)
+            if type(value) is dict:
+                inner = ET.Element(key, value)
                 ctl.append(inner)
+            else:
+                ctl.set(key, value)
         return ctl
 
     def __str__(self, *args, **kwargs):
@@ -298,18 +332,44 @@ class VacBotCommand:
 
 
 class Clean(VacBotCommand):
-    def __init__(self, wait):
+    def __init__(self, mode='auto', speed='normal', wait=None, terminal=False):
-        super().__init__('Clean', {'type': 'auto', 'speed': 'standard'}, wait)
+        super().__init__('Clean', {'clean': {'type': self.CLEAN_MODE[mode], 'speed': self.FAN_SPEED[speed]}}, wait=wait, terminal=terminal )
 
 
-class Edge(VacBotCommand):
+class Edge(Clean):
-    def __init__(self, wait):
+    def __init__(self, wait=None, terminal=False):
-        super().__init__('Clean', {'type': 'border', 'speed': 'strong'}, wait)
+        super().__init__('edge', 'high', wait=wait, terminal=terminal)
+
+
+class Spot(Clean):
+    def __init__(self, wait=None, terminal=False):
+        super().__init__('spot', 'high', wait=wait, terminal=terminal)
+
+
+class Stop(Clean):
+    def __init__(self, wait=None, terminal=False):
+        super().__init__('stop', 'normal', wait=wait, terminal=terminal)
+
+
+class StopAndWaitForCompletion(Stop):
+    def __init__(self, terminal=False):
+        super().__init__(terminal=False)
+
+    def wait_for_completion(self, bot):
+        logging.debug("waiting in " + self.name)
+        while bot.clean_status not in ['stop']:
+            time.sleep(0.5)
+        logging.debug("done waiting in " + self.name)
 
 
 class Charge(VacBotCommand):
-    def __init__(self):
+    def __init__(self, terminal=False):
-        super().__init__('Charge', {'type': 'go'}, terminal=True)
+        super().__init__('Charge', {'charge': {'type': self.CHARGE_MODE['return']}}, terminal=terminal)
+
+
+class ChargeAndWaitForCompletion(VacBotCommand):
+    def __init__(self, terminal=False):
+        super().__init__(terminal=terminal)
 
     def wait_for_completion(self, bot):
         logging.debug("waiting in " + self.name)
@@ -319,12 +379,31 @@ class Charge(VacBotCommand):
         click.echo("docked")
 
 
-class Stop(VacBotCommand):
+class Move(VacBotCommand):
-    def __init__(self):
+    def __init__(self, action):
-        super().__init__('Clean', {'type': 'stop', 'speed': 'standard'}, terminal=True)
+        super().__init__('Move', {'move': {'action': self.ACTION[action]}})
 
-    def wait_for_completion(self, bot):
+
-        logging.debug("waiting in " + self.name)
+class GetCleanState(VacBotCommand):
-        while bot.clean_status not in ['stop']:
+    def __init__(self):
-            time.sleep(0.5)
+        super().__init__('GetCleanState')
-        logging.debug("done waiting in " + self.name)
+
+
+class GetChargeState(VacBotCommand):
+    def __init__(self):
+        super().__init__('GetChargeState')
+
+
+class GetBatteryState(VacBotCommand):
+    def __init__(self):
+        super().__init__('GetBatteryInfo')
+
+
+class GetLifeSpan(VacBotCommand):
+    def __init__(self, component):
+        super().__init__('GetLifeSpan', {'type':self.COMPONENT[component]})
+
+
+class SetTime(VacBotCommand):
+    def __init__(self, timestamp, timezone):
+        super().__init__('SetTime', {'time':{'t':timestamp, 'tz':timezone}})

sucks/cli.py

@@ -125,7 +125,7 @@ def login(email, password, country_code, continent_code):
 @click.argument('minutes', type=click.FLOAT)
 def clean(frequency, minutes):
     if should_run(frequency):
-        return Clean(minutes * 60)
+        return Clean(wait=minutes * 60)
 
 
 @cli.command(help='cleans room edges for the specified number of minutes')
@@ -133,17 +133,17 @@ def clean(frequency, minutes):
 @click.argument('minutes', type=click.FLOAT)
 def edge(frequency, minutes):
     if should_run(frequency):
-        return Edge(minutes * 60)
+        return Edge(wait=minutes * 60)
 
 
 @cli.command(help='returns to charger')
 def charge():
-    return Charge()
+    return ChargeAndWaitForCompletion(terminal=True)
 
 
 @cli.command(help='stops the robot in its current position')
 def stop():
-    return Stop()
+    return StopAndWaitForCompletion(terminal=True)
 
 
 @cli.resultcallback()

tests/test_commands.py

@@ -9,8 +9,14 @@ 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'<ctl td="CustomCommand" type="customtype" />')
 
-                  b'<ctl td="CustomCommand"><customcommand type="customtype" /></ctl>')
+
+def test_custom_command_inner_tag():
+    # Ensure a custom-built command generates the expected XML payload
+    c = VacBotCommand('CustomCommand', {'customtag': {'customvar': 'customvalue'}})
+    assert_equals(ElementTree.tostring(c.to_xml()),
+                  b'<ctl td="CustomCommand"><customtag customvar="customvalue" /></ctl>')
 
 
 def test_custom_command_noargs():
@@ -21,31 +27,96 @@ def test_custom_command_noargs():
 
 
 def test_clean_command():
-    c = Clean(10)
+    c = Clean()
     assert_equals(c.terminal, False)
-    assert_equals(c.wait, 10)
     assert_equals(ElementTree.tostring(c.to_xml()),
                   b'<ctl td="Clean"><clean speed="standard" type="auto" /></ctl>')  # protocol has attribs in other order
+    c = Clean('edge', 'high', 10, terminal=True)
+    assert_equals(c.wait, 10)
+    assert_equals(c.terminal, True)
+    assert_equals(ElementTree.tostring(c.to_xml()),
+                  b'<ctl td="Clean"><clean speed="strong" type="border" /></ctl>')  # 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(wait=10)
-    c = Edge(10)
     assert_equals(c.terminal, False)
     assert_equals(c.wait, 10)
     assert_equals(ElementTree.tostring(c.to_xml()),
                   b'<ctl td="Clean"><clean speed="strong" type="border" /></ctl>')  # protocol has attribs in other order
 
 
+def test_spot_command():
+    c = Spot(wait=5)
+    assert_equals(c.terminal, False)
+    assert_equals(c.wait, 5)
+    assert_equals(ElementTree.tostring(c.to_xml()),
+                  b'<ctl td="Clean"><clean speed="strong" type="spot" /></ctl>')  # protocol has attribs in other order
+
+
 def test_charge_command():
-    c = Charge()
+    c = Charge(terminal=True)
     assert_equals(c.terminal, True)
     assert_equals(ElementTree.tostring(c.to_xml()),
                   b'<ctl td="Charge"><charge type="go" /></ctl>')
 
 
 def test_stop_command():
-    c = Stop()
+    c = Stop(terminal=True)
     assert_equals(c.terminal, True)
     assert_equals(ElementTree.tostring(c.to_xml()),
                   b'<ctl td="Clean"><clean speed="standard" type="stop" /></ctl>')
+
+
+def test_get_clean_state_command():
+    c = GetCleanState()
+    assert_equals(c.terminal, False)
+    assert_equals(ElementTree.tostring(c.to_xml()),
+                  b'<ctl td="GetCleanState" />')
+
+
+def test_get_charge_state_command():
+    c = GetChargeState()
+    assert_equals(c.terminal, False)
+    assert_equals(ElementTree.tostring(c.to_xml()),
+                  b'<ctl td="GetChargeState" />')
+
+
+def test_get_battery_state_command():
+    c = GetBatteryState()
+    assert_equals(c.terminal, False)
+    assert_equals(ElementTree.tostring(c.to_xml()),
+                  b'<ctl td="GetBatteryInfo" />')
+
+
+
+def test_move_command():
+    c = Move(action='left')
+    assert_equals(c.terminal, False)
+    assert_equals(ElementTree.tostring(c.to_xml()),
+                  b'<ctl td="Move"><move action="SpinLeft" /></ctl>')
+    c = Move(action='right')
+    assert_equals(ElementTree.tostring(c.to_xml()),
+                  b'<ctl td="Move"><move action="SpinRight" /></ctl>')
+    c = Move(action='turn_around')
+    assert_equals(ElementTree.tostring(c.to_xml()),
+                  b'<ctl td="Move"><move action="TurnAround" /></ctl>')
+    c = Move(action='forward')
+    assert_equals(ElementTree.tostring(c.to_xml()),
+                  b'<ctl td="Move"><move action="forward" /></ctl>')
+    c = Move(action='stop')
+    assert_equals(ElementTree.tostring(c.to_xml()),
+                  b'<ctl td="Move"><move action="stop" /></ctl>')
+
+
+def test_get_lifepsan_command():
+    c = GetLifeSpan('main_brush')
+    assert_equals(c.terminal, False)
+    assert_equals(ElementTree.tostring(c.to_xml()),
+                  b'<ctl td="GetLifeSpan" type="Brush" />')
+    c = GetLifeSpan('side_brush')
+    assert_equals(ElementTree.tostring(c.to_xml()),
+                  b'<ctl td="GetLifeSpan" type="SideBrush" />')
+    c = GetLifeSpan('filter')
+    assert_equals(ElementTree.tostring(c.to_xml()),
+                  b'<ctl td="GetLifeSpan" type="DustCaseHeap" />')

tests/test_ecovacs_xmpp.py

@@ -10,7 +10,7 @@ from sucks import *
 
 def test_wrap_command():
     x = make_ecovacs_xmpp()
-    c = str(x._wrap_command(Clean(1).to_xml(), 'E0000000001234567890@126.ecorobot.net/atom'))
+    c = str(x._wrap_command(Clean(wait=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))