python-kasa · rytilahti · Apr 23, 2024 · Apr 22, 2024 · Apr 22, 2024 · Apr 22, 2024
diff --git a/kasa/smart/modules/cloudmodule.py b/kasa/smart/modules/cloudmodule.py
@@ -4,6 +4,7 @@
 
 from typing import TYPE_CHECKING
 
+from ...exceptions import SmartErrorCode
 from ...feature import Feature, FeatureType
 from ..smartmodule import SmartModule
 
@@ -34,4 +35,6 @@ def __init__(self, device: SmartDevice, module: str):
     @property
     def is_connected(self):
         """Return True if device is connected to the cloud."""
+        if isinstance(self.data, SmartErrorCode):
+            return False
         return self.data["status"] == 0
diff --git a/kasa/smart/modules/firmware.py b/kasa/smart/modules/firmware.py
@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Optional
+from typing import TYPE_CHECKING, Any, Optional
 
 from ...exceptions import SmartErrorCode
 from ...feature import Feature, FeatureType
@@ -74,9 +74,7 @@ def __init__(self, device: SmartDevice, module: str):
 
     def query(self) -> dict:
         """Query to execute during the update cycle."""
-        req = {
-            "get_latest_fw": None,
-        }
+        req: dict[str, Any] = {"get_latest_fw": None}
         if self.supported_version > 1:
             req["get_auto_update_info"] = None
         return req
@@ -85,15 +83,17 @@ def query(self) -> dict:
     def latest_firmware(self):
         """Return latest firmware information."""
         fw = self.data.get("get_latest_fw") or self.data
-        if isinstance(fw, SmartErrorCode):
+        if not self._device.is_cloud_connected or isinstance(fw, SmartErrorCode):
             # Error in response, probably disconnected from the cloud.
             return UpdateInfo(type=0, need_to_upgrade=False)
 
         return UpdateInfo.parse_obj(fw)
 
     @property
-    def update_available(self):
+    def update_available(self) -> bool | None:
         """Return True if update is available."""
+        if not self._device.is_cloud_connected:
+            return None
         return self.latest_firmware.update_available
 
     async def get_update_state(self):

diff --git a/kasa/smart/smartdevice.py b/kasa/smart/smartdevice.py
@@ -104,7 +104,11 @@ async def _negotiate(self):
         We fetch the device info and the available components as early as possible.
         If the device reports supporting child devices, they are also initialized.
         """
-        initial_query = {"component_nego": None, "get_device_info": None}
+        initial_query = {
+            "component_nego": None,
+            "get_device_info": None,
+            "get_connect_cloud_state": None,
+        }
         resp = await self.protocol.query(initial_query)
 
         # Save the initial state to allow modules access the device info already
@@ -238,6 +242,13 @@ async def _initialize_features(self):
             for feat in module._module_features.values():
                 self._add_feature(feat)
 
+    @property
+    def is_cloud_connected(self):
+        """Returns if the device is connected to the cloud."""
+        if "CloudModule" not in self.modules:
+            return False
+        return self.modules["CloudModule"].is_connected
+
     @property
     def sys_info(self) -> dict[str, Any]:
         """Returns the device info."""

diff --git a/kasa/tests/fakeprotocol_smart.py b/kasa/tests/fakeprotocol_smart.py
@@ -65,7 +65,6 @@ def credentials_hash(self):
                 },
             },
         ),
-        "get_connect_cloud_state": ("cloud_connect", {"status": 1}),
         "get_on_off_gradually_info": ("on_off_gradually", {"enable": True}),
         "get_latest_fw": (
             "firmware",
@@ -172,7 +171,7 @@ def _send_request(self, request_dict: dict):
                 # calling the unsupported device in the first place.
                 retval = {
                     "error_code": SmartErrorCode.PARAMS_ERROR.value,
-                    "method": "get_device_usage",
+                    "method": method,
                 }
             # Reduce warning spam by consolidating and reporting at the end of the run
             if self.fixture_name not in pytest.fixtures_missing_methods:

diff --git a/kasa/tests/test_smartdevice.py b/kasa/tests/test_smartdevice.py
@@ -4,6 +4,7 @@
 
 import logging
 from typing import Any
+from unittest.mock import patch
 
 import pytest
 from pytest_mock import MockerFixture
@@ -77,7 +78,13 @@ async def test_negotiate(dev: SmartDevice, mocker: MockerFixture):
     await dev._negotiate()
 
     # Check that we got the initial negotiation call
-    query.assert_any_call({"component_nego": None, "get_device_info": None})
+    query.assert_any_call(
+        {
+            "component_nego": None,
+            "get_device_info": None,
+            "get_connect_cloud_state": None,
+        }
+    )
     assert dev._components_raw
 
     # Check the children are created, if device supports them
@@ -128,3 +135,62 @@ async def test_smartdevice_brightness(dev: SmartBulb):
 
     with pytest.raises(ValueError):
         await dev.set_brightness(feature.maximum_value + 10)
+
+
+@device_smart
+async def test_smartdevice_cloud_connection(dev: SmartDevice, mocker: MockerFixture):
+    """Test is_cloud_connected property."""
+    assert isinstance(dev, SmartDevice)
+    assert "cloud_connect" in dev._components
+
+    is_connected = (
+        (cc := dev._last_update.get("get_connect_cloud_state"))
+        and not isinstance(cc, SmartErrorCode)
+        and cc["status"] == 0
+    )
+
+    assert dev.is_cloud_connected == is_connected
+    last_update = dev._last_update
+
+    last_update["get_connect_cloud_state"] = {"status": 0}
+    with patch.object(dev.protocol, "query", return_value=last_update):
+        await dev.update()
+        assert dev.is_cloud_connected is True
+
+    last_update["get_connect_cloud_state"] = {"status": 1}
+    with patch.object(dev.protocol, "query", return_value=last_update):
+        await dev.update()
+        assert dev.is_cloud_connected is False
+
+    last_update["get_connect_cloud_state"] = SmartErrorCode.UNKNOWN_METHOD_ERROR
+    with patch.object(dev.protocol, "query", return_value=last_update):
+        await dev.update()
+        assert dev.is_cloud_connected is False
+
+    # Test for no cloud_connect component during device initialisation
+    component_list = [
+        val
+        for val in dev._components_raw["component_list"]
+        if val["id"] not in {"cloud_connect"}
+    ]
+    initial_response = {
+        "component_nego": {"component_list": component_list},
+        "get_connect_cloud_state": last_update["get_connect_cloud_state"],
+        "get_device_info": last_update["get_device_info"],
+    }
+    # Child component list is not stored on the device
+    if "get_child_device_list" in last_update:
+        child_component_list = await dev.protocol.query(
+            "get_child_device_component_list"
+        )
+        last_update["get_child_device_component_list"] = child_component_list[
+            "get_child_device_component_list"
+        ]
+    new_dev = SmartDevice("127.0.0.1", protocol=dev.protocol)
+    with patch.object(
+        new_dev.protocol,
+        "query",
+        side_effect=[initial_response, last_update, last_update],
+    ):
+        await new_dev.update()
+        assert new_dev.is_cloud_connected is False