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):