From 157ad8e8071c2ca2fd4c3584660b03b12c445147 Mon Sep 17 00:00:00 2001
From: "Steven B." <51370195+sdb9696@users.noreply.github.com>
Date: Wed, 13 Nov 2024 14:57:42 +0000
Subject: [PATCH] Update cli energy command to use energy module (#1252)

---
 kasa/cli/usage.py | 33 ++++++++++++++-------------------
 tests/test_cli.py | 19 +++++++++++--------
 2 files changed, 25 insertions(+), 27 deletions(-)

diff --git a/kasa/cli/usage.py b/kasa/cli/usage.py
index 314182fd..90a0fa78 100644
--- a/kasa/cli/usage.py
+++ b/kasa/cli/usage.py
@@ -9,11 +9,9 @@ import asyncclick as click
 
 from kasa import (
     Device,
+    Module,
 )
-from kasa.iot import (
-    IotDevice,
-)
-from kasa.iot.iotstrip import IotStripPlug
+from kasa.interfaces import Energy
 from kasa.iot.modules import Usage
 
 from .common import (
@@ -49,42 +47,39 @@ async def energy(dev: Device, year, month, erase):
     Daily and monthly data provided in CSV format.
     """
     echo("[bold]== Emeter ==[/bold]")
-    if not dev.has_emeter:
-        error("Device has no emeter")
+    if not (energy := dev.modules.get(Module.Energy)):
+        error("Device has no energy module.")
         return
 
-    if (year or month or erase) and not isinstance(dev, IotDevice):
-        error("Device has no historical statistics")
+    if (year or month or erase) and not energy.supports(
+        Energy.ModuleFeature.PERIODIC_STATS
+    ):
+        error("Device does not support historical statistics")
         return
-    else:
-        dev = cast(IotDevice, dev)
 
     if erase:
         echo("Erasing emeter statistics..")
-        return await dev.erase_emeter_stats()
+        return await energy.erase_stats()
 
     if year:
         echo(f"== For year {year.year} ==")
         echo("Month, usage (kWh)")
-        usage_data = await dev.get_emeter_monthly(year=year.year)
+        usage_data = await energy.get_monthly_stats(year=year.year)
     elif month:
         echo(f"== For month {month.month} of {month.year} ==")
         echo("Day, usage (kWh)")
-        usage_data = await dev.get_emeter_daily(year=month.year, month=month.month)
+        usage_data = await energy.get_daily_stats(year=month.year, month=month.month)
     else:
         # Call with no argument outputs summary data and returns
-        if isinstance(dev, IotStripPlug):
-            emeter_status = await dev.get_emeter_realtime()
-        else:
-            emeter_status = dev.emeter_realtime
+        emeter_status = await energy.get_status()
 
         echo("Current: {} A".format(emeter_status["current"]))
         echo("Voltage: {} V".format(emeter_status["voltage"]))
         echo("Power: {} W".format(emeter_status["power"]))
         echo("Total consumption: {} kWh".format(emeter_status["total"]))
 
-        echo(f"Today: {dev.emeter_today} kWh")
-        echo(f"This month: {dev.emeter_this_month} kWh")
+        echo(f"Today: {energy.consumption_today} kWh")
+        echo(f"This month: {energy.consumption_this_month} kWh")
 
         return emeter_status
 
diff --git a/tests/test_cli.py b/tests/test_cli.py
index b6bcdfd4..24fc916f 100644
--- a/tests/test_cli.py
+++ b/tests/test_cli.py
@@ -15,6 +15,7 @@ from kasa import (
     Credentials,
     Device,
     DeviceError,
+    DeviceType,
     EmeterStatus,
     KasaException,
     Module,
@@ -424,20 +425,22 @@ async def test_time_set(dev: Device, mocker, runner):
 
 async def test_emeter(dev: Device, mocker, runner):
     res = await runner.invoke(emeter, obj=dev)
-    if not dev.has_emeter:
-        assert "Device has no emeter" in res.output
+    if not (energy := dev.modules.get(Module.Energy)):
+        assert "Device has no energy module." in res.output
         return
 
     assert "== Emeter ==" in res.output
 
-    if not dev.is_strip:
+    if dev.device_type is not DeviceType.Strip:
         res = await runner.invoke(emeter, ["--index", "0"], obj=dev)
         assert f"Device: {dev.host} does not have children" in res.output
         res = await runner.invoke(emeter, ["--name", "mock"], obj=dev)
         assert f"Device: {dev.host} does not have children" in res.output
 
-    if dev.is_strip and len(dev.children) > 0:
-        realtime_emeter = mocker.patch.object(dev.children[0], "get_emeter_realtime")
+    if dev.device_type is DeviceType.Strip and len(dev.children) > 0:
+        child_energy = dev.children[0].modules.get(Module.Energy)
+        assert child_energy
+        realtime_emeter = mocker.patch.object(child_energy, "get_status")
         realtime_emeter.return_value = EmeterStatus({"voltage_mv": 122066})
 
         res = await runner.invoke(emeter, ["--index", "0"], obj=dev)
@@ -450,18 +453,18 @@ async def test_emeter(dev: Device, mocker, runner):
         assert realtime_emeter.call_count == 2
 
     if isinstance(dev, IotDevice):
-        monthly = mocker.patch.object(dev, "get_emeter_monthly")
+        monthly = mocker.patch.object(energy, "get_monthly_stats")
         monthly.return_value = {1: 1234}
     res = await runner.invoke(emeter, ["--year", "1900"], obj=dev)
     if not isinstance(dev, IotDevice):
-        assert "Device has no historical statistics" in res.output
+        assert "Device does not support historical statistics" in res.output
         return
     assert "For year" in res.output
     assert "1, 1234" in res.output
     monthly.assert_called_with(year=1900)
 
     if isinstance(dev, IotDevice):
-        daily = mocker.patch.object(dev, "get_emeter_daily")
+        daily = mocker.patch.object(energy, "get_daily_stats")
         daily.return_value = {1: 1234}
     res = await runner.invoke(emeter, ["--month", "1900-12"], obj=dev)
     if not isinstance(dev, IotDevice):