Make get_module return typed module (#892)

Passing in a string still works and returns either `IotModule` or
`SmartModule` type when called on `IotDevice` or `SmartDevice`
respectively. When calling on `Device` will return `Module` type.

Passing in a module type is then typed to that module, i.e.:
```py
smartdev.get_module(FanModule)  # type is FanModule
smartdev.get_module("FanModule")  # type is SmartModule
```
Only thing this doesn't do is check that you can't pass an `IotModule`
to a `SmartDevice.get_module()`. However there is a runtime check which
will return null if the passed `ModuleType` is not a subclass of
`SmartModule`.

Many thanks to @cdce8p for helping with this.

kasa/tests/smart/features/test_brightness.py

@@ -33,7 +33,7 @@ async def test_brightness_component(dev: SmartDevice):
 
 
 @dimmable
-async def test_brightness_dimmable(dev: SmartDevice):
+async def test_brightness_dimmable(dev: IotDevice):
     """Test brightness feature."""
     assert isinstance(dev, IotDevice)
     assert "brightness" in dev.sys_info or bool(dev.sys_info["is_dimmable"])

kasa/tests/smart/modules/test_fan.py

@@ -1,5 +1,3 @@
-from typing import cast
-
 import pytest
 from pytest_mock import MockerFixture
 
@@ -13,7 +11,7 @@ fan = parametrize("has fan", component_filter="fan_control", protocol_filter={"S
 @fan
 async def test_fan_speed(dev: SmartDevice, mocker: MockerFixture):
     """Test fan speed feature."""
-    fan = cast(FanModule, dev.get_module("FanModule"))
+    fan = dev.get_module(FanModule)
     assert fan
 
     level_feature = fan._module_features["fan_speed_level"]
@@ -38,7 +36,7 @@ async def test_fan_speed(dev: SmartDevice, mocker: MockerFixture):
 @fan
 async def test_sleep_mode(dev: SmartDevice, mocker: MockerFixture):
     """Test sleep mode feature."""
-    fan = cast(FanModule, dev.get_module("FanModule"))
+    fan = dev.get_module(FanModule)
     assert fan
     sleep_feature = fan._module_features["fan_sleep_mode"]
     assert isinstance(sleep_feature.value, bool)
@@ -57,7 +55,8 @@ async def test_sleep_mode(dev: SmartDevice, mocker: MockerFixture):
 async def test_fan_interface(dev: SmartDevice, mocker: MockerFixture):
     """Test fan speed on device interface."""
     assert isinstance(dev, SmartDevice)
-    fan = cast(FanModule, dev.get_module("FanModule"))
+    fan = dev.get_module(FanModule)
+    assert fan
     device = fan._device
     assert device.is_fan
 

kasa/tests/test_iotdevice.py

@@ -19,7 +19,7 @@ from voluptuous import (
 from kasa import KasaException
 from kasa.iot import IotDevice
 
-from .conftest import handle_turn_on, turn_on
+from .conftest import get_device_for_fixture_protocol, handle_turn_on, turn_on
 from .device_fixtures import device_iot, has_emeter_iot, no_emeter_iot
 from .fakeprotocol_iot import FakeIotProtocol
 
@@ -258,3 +258,30 @@ async def test_modules_not_supported(dev: IotDevice):
     await dev.update()
     for module in dev.modules.values():
         assert module.is_supported is not None
+
+
+async def test_get_modules():
+    """Test get_modules for child and parent modules."""
+    dummy_device = await get_device_for_fixture_protocol(
+        "HS100(US)_2.0_1.5.6.json", "IOT"
+    )
+    from kasa.iot.modules import Cloud
+    from kasa.smart.modules import CloudModule
+
+    # Modules on device
+    module = dummy_device.get_module("Cloud")
+    assert module
+    assert module._device == dummy_device
+    assert isinstance(module, Cloud)
+
+    module = dummy_device.get_module(Cloud)
+    assert module
+    assert module._device == dummy_device
+    assert isinstance(module, Cloud)
+
+    # Invalid modules
+    module = dummy_device.get_module("DummyModule")
+    assert module is None
+
+    module = dummy_device.get_module(CloudModule)
+    assert module is None

kasa/tests/test_smartdevice.py

@@ -122,23 +122,43 @@ async def test_update_module_queries(dev: SmartDevice, mocker: MockerFixture):
             spies[device].assert_not_called()
 
 
-async def test_get_modules(mocker):
+async def test_get_modules():
     """Test get_modules for child and parent modules."""
     dummy_device = await get_device_for_fixture_protocol(
         "KS240(US)_1.0_1.0.5.json", "SMART"
     )
+    from kasa.iot.modules import AmbientLight
+    from kasa.smart.modules import CloudModule, FanModule
+
+    # Modules on device
     module = dummy_device.get_module("CloudModule")
     assert module
     assert module._device == dummy_device
+    assert isinstance(module, CloudModule)
 
+    module = dummy_device.get_module(CloudModule)
+    assert module
+    assert module._device == dummy_device
+    assert isinstance(module, CloudModule)
+
+    # Modules on child
     module = dummy_device.get_module("FanModule")
     assert module
     assert module._device != dummy_device
     assert module._device._parent == dummy_device
 
+    module = dummy_device.get_module(FanModule)
+    assert module
+    assert module._device != dummy_device
+    assert module._device._parent == dummy_device
+
+    # Invalid modules
     module = dummy_device.get_module("DummyModule")
     assert module is None
 
+    module = dummy_device.get_module(AmbientLight)
+    assert module is None
+
 
 @bulb_smart
 async def test_smartdevice_brightness(dev: SmartDevice):