From e576fcdb463b602d4ae9042a02cc36a32902be26 Mon Sep 17 00:00:00 2001
From: Teemu R <tpr@iki.fi>
Date: Tue, 23 Jan 2024 22:58:57 +0100
Subject: [PATCH] Allow raw-command and wifi without update (#688)

* Allow raw-command and wifi without update

* Call update always but on wifi&raw-command

* Add tests

* Skip update also if device_family was defined, as device factory performs an update
---
 kasa/cli.py            |  4 +++-
 kasa/tests/test_cli.py | 35 +++++++++++++++++++++++++++++++++--
 2 files changed, 36 insertions(+), 3 deletions(-)

diff --git a/kasa/cli.py b/kasa/cli.py
index d1cb7276..1b20303d 100755
--- a/kasa/cli.py
+++ b/kasa/cli.py
@@ -317,7 +317,6 @@ async def cli(
 
     if type is not None:
         dev = TYPE_TO_CLASS[type](host)
-        await dev.update()
     elif device_family and encrypt_type:
         ctype = ConnectionType(
             DeviceFamilyType(device_family),
@@ -339,6 +338,9 @@ async def cli(
             port=port,
             credentials=credentials,
         )
+
+    # Skip update for wifi & raw-command, and if factory was used to connect
+    if ctx.invoked_subcommand not in ["wifi", "raw-command"] and not device_family:
         await dev.update()
 
     ctx.obj = dev
diff --git a/kasa/tests/test_cli.py b/kasa/tests/test_cli.py
index 3aad37dd..fa2d5c69 100644
--- a/kasa/tests/test_cli.py
+++ b/kasa/tests/test_cli.py
@@ -34,6 +34,27 @@ from kasa.smartprotocol import SmartProtocol
 from .conftest import device_iot, device_smart, handle_turn_on, new_discovery, turn_on
 
 
+async def test_update_called_by_cli(dev, mocker):
+    """Test that device update is called on main."""
+    runner = CliRunner()
+    update = mocker.patch.object(dev, "update")
+    mocker.patch("kasa.discover.Discover.discover_single", return_value=dev)
+
+    res = await runner.invoke(
+        cli,
+        [
+            "--host",
+            "127.0.0.1",
+            "--username",
+            "foo",
+            "--password",
+            "bar",
+        ],
+    )
+    assert res.exit_code == 0
+    update.assert_called()
+
+
 @device_iot
 async def test_sysinfo(dev):
     runner = CliRunner()
@@ -86,8 +107,9 @@ async def test_alias(dev):
     await dev.set_alias(old_alias)
 
 
-async def test_raw_command(dev):
+async def test_raw_command(dev, mocker):
     runner = CliRunner()
+    update = mocker.patch.object(dev, "update")
     from kasa.tapo import TapoDevice
 
     if isinstance(dev, TapoDevice):
@@ -96,6 +118,10 @@ async def test_raw_command(dev):
         params = ["system", "get_sysinfo"]
     res = await runner.invoke(raw_command, params, obj=dev)
 
+    # Make sure that update was not called for wifi
+    with pytest.raises(AssertionError):
+        update.assert_called()
+
     assert res.exit_code == 0
     assert dev.model in res.output
 
@@ -129,14 +155,19 @@ async def test_wifi_scan(dev):
 
 
 @device_smart
-async def test_wifi_join(dev):
+async def test_wifi_join(dev, mocker):
     runner = CliRunner()
+    update = mocker.patch.object(dev, "update")
     res = await runner.invoke(
         wifi,
         ["join", "FOOBAR", "--keytype", "wpa_psk", "--password", "foobar"],
         obj=dev,
     )
 
+    # Make sure that update was not called for wifi
+    with pytest.raises(AssertionError):
+        update.assert_called()
+
     assert res.exit_code == 0
     assert "Asking the device to connect to FOOBAR" in res.output