diff --git a/doc/configuration.rst b/doc/configuration.rst index 60bfbb494..dcbd04108 100644 --- a/doc/configuration.rst +++ b/doc/configuration.rst @@ -249,6 +249,14 @@ Currently available are: `__ for details. +``tapo`` + Controls *Tapo power strips and single socket devices* via `python-kasa + `_. + Requires valid TP-Link/TAPO cloud credentials to work. + See the `docstring in the module + `__ + for details. + ``tinycontrol`` Controls a tinycontrol.eu IP Power Socket via HTTP. It was tested on the *6G10A v2* model. diff --git a/labgrid/driver/power/tapo.py b/labgrid/driver/power/tapo.py new file mode 100644 index 000000000..590b5408e --- /dev/null +++ b/labgrid/driver/power/tapo.py @@ -0,0 +1,106 @@ +"""Driver for controlling TP-Link Tapo smart plugs and power strips. + +This module provides functionality to control TP-Link Tapo smart power devices through +the kasa library. It supports both single socket devices (like P100) and multi-socket +power strips (like P300). + +Features: +- Environment-based authentication using KASA_USERNAME and KASA_PASSWORD +- Support for both single and multi-socket devices + +Requirements: +- Valid TP-Link cloud credentials (username/password) +""" + +import asyncio +import os + +import kasa +from kasa import Credentials, Device, DeviceConfig, DeviceConnectionParameters, DeviceEncryptionType, DeviceFamily +from packaging.version import parse + + +def _using_old_kasa_api() -> bool: + target_kasa_version = parse("0.10.0") + current_kasa_version = parse(kasa.__version__) + return current_kasa_version < target_kasa_version + + +def _get_credentials() -> Credentials: + username = os.environ.get("KASA_USERNAME") + password = os.environ.get("KASA_PASSWORD") + if username is None or password is None: + raise EnvironmentError("KASA_USERNAME or KASA_PASSWORD environment variable not set") + return Credentials(username=username, password=password) + + +def _get_connection_type() -> DeviceConnectionParameters: + # API changed between versions 0.9.1 and 0.10.0 of python-kasa + if _using_old_kasa_api(): + return DeviceConnectionParameters( + device_family=DeviceFamily.SmartTapoPlug, + encryption_type=DeviceEncryptionType.Klap, + https=False, + login_version=2, + ) + return DeviceConnectionParameters( + device_family=DeviceFamily.SmartTapoPlug, + encryption_type=DeviceEncryptionType.Klap, + https=False, + login_version=2, + http_port=80, + ) + + +def _get_device_config(host: str) -> DeviceConfig: + # Same as with `_get_connection_type` - python-kasa API changed + if _using_old_kasa_api(): + return DeviceConfig( + host=host, credentials=_get_credentials(), connection_type=_get_connection_type(), uses_http=True + ) # type: ignore[call-arg] + return DeviceConfig( + host=host, credentials=_get_credentials(), connection_type=_get_connection_type() + ) + + +async def _power_set(host: str, port: str, index: str, value: bool) -> None: + """We embed the coroutines in an `async` function to minimise calls to `asyncio.run`""" + assert port is None + index = int(index) + device = await Device.connect(config=_get_device_config(host)) + await device.update() + + if device.children: + assert len(device.children) > index, "Trying to access non-existant plug socket on device" + + target = device if not device.children else device.children[index] + if value: + await target.turn_on() + else: + await target.turn_off() + await device.disconnect() + + +def power_set(host: str, port: str, index: str, value: bool) -> None: + asyncio.run(_power_set(host, port, index, value)) + + +async def _power_get(host: str, port: str, index: str) -> bool: + assert port is None + index = int(index) + device = await Device.connect(config=_get_device_config(host)) + await device.update() + + pwr_state: bool + # If the device has no children, it is a single plug socket + if not device.children: + pwr_state = device.is_on + else: + assert len(device.children) > index, "Trying to access non-existant plug socket on device" + pwr_state = device.children[index].is_on + await device.disconnect() + return pwr_state + + +def power_get(host: str, port: str, index: str) -> bool: + return asyncio.run(_power_get(host, port, index)) diff --git a/pyproject.toml b/pyproject.toml index 33caa6f31..ac5c14c4e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -80,6 +80,7 @@ dev = [ # additional dev dependencies "psutil>=5.8.0", + "pytest-asyncio>=0.25.3", "pytest-benchmark>=4.0.0", "pytest-cov>=3.0.0", "pytest-dependency>=0.5.1", diff --git a/tests/test_powerdriver.py b/tests/test_powerdriver.py index 2ae8783b2..1093e2f4c 100644 --- a/tests/test_powerdriver.py +++ b/tests/test_powerdriver.py @@ -300,6 +300,10 @@ def test_import_backend_tplink(self): pytest.importorskip("kasa") import labgrid.driver.power.tplink + def test_import_backend_tapo(self): + pytest.importorskip("kasa") + import labgrid.driver.power.tapo + def test_import_backend_siglent(self): pytest.importorskip("vxi11") import labgrid.driver.power.siglent diff --git a/tests/test_tapo.py b/tests/test_tapo.py new file mode 100644 index 000000000..652c1414b --- /dev/null +++ b/tests/test_tapo.py @@ -0,0 +1,142 @@ +import os +from unittest.mock import AsyncMock, patch + +import pytest + +from labgrid.driver.power.tapo import _get_credentials, _power_get, _power_set, power_get + + +@pytest.fixture +def mock_device_strip(): + device = AsyncMock() + device.children = [ + AsyncMock(is_on=True), + AsyncMock(is_on=False), + AsyncMock(is_on=True) + ] + return device + + +@pytest.fixture +def mock_device_single_plug(): + device = AsyncMock() + device.children = [] + return device + + +@pytest.fixture +def mock_env(): + os.environ['KASA_USERNAME'] = 'test_user' + os.environ['KASA_PASSWORD'] = 'test_pass' + yield + del os.environ['KASA_USERNAME'] + del os.environ['KASA_PASSWORD'] + + +class TestTapoPowerDriver: + def test_get_credentials_should_raise_value_error_when_credentials_missing(self): + # Save existing environment variables + saved_username = os.environ.pop('KASA_USERNAME', None) + saved_password = os.environ.pop('KASA_PASSWORD', None) + + try: + with pytest.raises(EnvironmentError, match="KASA_USERNAME or KASA_PASSWORD environment variable not set"): + _get_credentials() + finally: + # Restore environment variables if they existed + if saved_username is not None: + os.environ['KASA_USERNAME'] = saved_username + if saved_password is not None: + os.environ['KASA_PASSWORD'] = saved_password + + def test_credentials_valid(self, mock_env): + creds = _get_credentials() + assert creds.username == 'test_user' + assert creds.password == 'test_pass' + + @pytest.mark.asyncio + async def test_power_get_single_plug_turn_on(self, mock_device_single_plug, mock_env): + mock_device_single_plug.is_on = True + + with patch('kasa.Device.connect', return_value=mock_device_single_plug): + result = await _power_get('192.168.1.100', None, "0") + assert result is True + + @pytest.mark.asyncio + async def test_power_get_single_plug_turn_off(self, mock_device_single_plug, mock_env): + mock_device_single_plug.is_on = False + + with patch('kasa.Device.connect', return_value=mock_device_single_plug): + result = await _power_get('192.168.1.100', None, "0") + assert result is False + + @pytest.mark.asyncio + async def test_power_get_single_plug_should_not_care_for_index(self, mock_device_single_plug, mock_env): + invalid_index_ignored = "7" + mock_device_single_plug.is_on = True + + with patch('kasa.Device.connect', return_value=mock_device_single_plug): + result = await _power_get('192.168.1.100', None, invalid_index_ignored) + assert result is True + + @pytest.mark.asyncio + async def test_power_set_single_plug_turn_on(self, mock_device_single_plug, mock_env): + mock_device_single_plug.is_on = False + with patch('kasa.Device.connect', return_value=mock_device_single_plug): + await _power_set('192.168.1.100', None, "0", True) + mock_device_single_plug.turn_on.assert_called_once() + + @pytest.mark.asyncio + async def test_power_set_single_plug_turn_off(self, mock_device_single_plug, mock_env): + mock_device_single_plug.is_on = True + with patch('kasa.Device.connect', return_value=mock_device_single_plug): + await _power_set('192.168.1.100', None, "0", False) + mock_device_single_plug.turn_off.assert_called_once() + + @pytest.mark.asyncio + async def test_power_get_strip_valid_socket(self, mock_device_strip, mock_env): + with patch('kasa.Device.connect', return_value=mock_device_strip): + # Test first outlet (on) + result = await _power_get('192.168.1.100', None, "0") + assert result is True + + # Test second outlet (off) + result = await _power_get('192.168.1.100', None, "1") + assert result is False + + # Test third outlet (on) + result = await _power_get('192.168.1.100', None, "2") + assert result is True + + @pytest.mark.asyncio + async def test_power_set_strip_valid_socket(self, mock_device_strip, mock_env): + with patch('kasa.Device.connect', return_value=mock_device_strip): + await _power_set('192.168.1.100', None, "0", False) + mock_device_strip.children[0].turn_off.assert_called_once() + + await _power_set('192.168.1.100', None, "1", True) + mock_device_strip.children[1].turn_on.assert_called_once() + + def test_power_get_should_raise_assertion_error_when_invalid_index_strip(self, mock_device_strip, mock_env): + invalid_socket = "5" + with patch('kasa.Device.connect', return_value=mock_device_strip): + with pytest.raises(AssertionError, match="Trying to access non-existant plug socket"): + power_get('192.168.1.100', None, invalid_socket) + + @pytest.mark.asyncio + async def test_power_set_should_raise_assertion_error_when_invalid_index_strip(self, mock_device_strip, mock_env): + invalid_socket = "5" + with patch('kasa.Device.connect', return_value=mock_device_strip): + with pytest.raises(AssertionError, match="Trying to access non-existant plug socket"): + await _power_set('192.168.1.100', None, invalid_socket, True) + + def test_port_not_none_strip(self, mock_device_strip): + with patch('kasa.Device.connect', return_value=mock_device_strip): + with pytest.raises(AssertionError): + power_get('192.168.1.100', '8080', "0") + + def test_port_not_none_single_socket(self, mock_device_single_plug): + mock_device_single_plug.is_on = True + with patch('kasa.Device.connect', return_value=mock_device_single_plug): + with pytest.raises(AssertionError): + power_get('192.168.1.100', '8080', "0")