Refactoring reset consumable

miio/vacuum.py

@@ -21,6 +21,12 @@ class TimerState(enum.Enum):
     On = "on"
     Off = "off"
 
+class Consumable(enum.Enum):
+    MainBrush = "main_brush_work_time"
+    SideBrush = "side_brush_work_time"
+    Filter = "filter_work_time"
+    SensorDirty = "sensor_dirty_time"
+
 class Vacuum(Device):
     """Main class representing the vacuum."""
 
@@ -110,10 +116,9 @@ def consumable_status(self) -> ConsumableStatus:
         """Return information about consumables."""
         return ConsumableStatus(self.send("get_consumable")[0])
 
-    def consumable_reset(self, name):
+    def consumable_reset(self, consumable: Consumable):
         """Reset consumable information."""
-        # name = ["main_brush_work_time", "side_brush_work_time", "sensor_dirty_time" or "filter_work_time"]
-        return self.send("reset_consumable", [name])
+        return self.send("reset_consumable", [consumable.value])
 
     def map(self):
         """Return map token."""

miio/vacuum_cli.py

@@ -148,9 +148,20 @@ def consumables(vac: miio.Vacuum):
 @click.argument('name', type=str, required=True)
 @pass_dev
 def reset_consumable(vac: miio.Vacuum, name):
-    """Query and reset consumable."""
-    click.echo("Reset consumable %s" % name)
-    vac.consumable_reset(name)
+    """Reset consumable."""
+    from miio.vacuum import Consumable
+    if name == 'main_brush':
+        consumable = Consumable.MainBrush
+    elif name == 'side_brush':
+        consumable = Consumable.SideBrush
+    elif name == 'filter':
+        consumable = Consumable.Filter
+    elif name == 'sensor_dirty':
+        consumable = Consumable.SensorDirty
+    else:
+        return
+
+    click.echo("Resetting consumable '%s': %s" % (name, vac.consumable_reset(consumable)))
 
 
 @cli.command()
@@ -454,6 +465,6 @@ def raw_command(vac: miio.Vacuum, cmd, parameters):
     click.echo("Sending cmd %s with params %s" % (cmd, params))
     click.echo(vac.raw_command(cmd, params))
 
-
+cli()
 if __name__ == "__main__":
     cli()