Skip to content

Commit d41c2fa

Browse files
MarekSzczypinskiMarekSzczypinski
committed
Add support for TAPO devices
Signed-off-by: Marek Szczypiński <[email protected]>
1 parent 2ce9e24 commit d41c2fa

File tree

5 files changed

+256
-0
lines changed

5 files changed

+256
-0
lines changed

‎doc/configuration.rst

+8
Original file line numberDiff line numberDiff line change
@@ -243,6 +243,14 @@ Currently available are:
243243
Controls *TP-Link power strips* via `python-kasa
244244
<https://github.com/python-kasa/python-kasa>`_.
245245

246+
``tapo``
247+
Controls *Tapo power strips and single socket devices* via `python-kasa
248+
<https://github.com/python-kasa/python-kasa>`_.
249+
Requires valid TP-Link/TAPO cloud credentials to work.
250+
See the `docstring in the module
251+
<https://github.com/labgrid-project/labgrid/blob/master/labgrid/driver/power/tapo.py>`__
252+
for details.
253+
246254
``tinycontrol``
247255
Controls a tinycontrol.eu IP Power Socket via HTTP.
248256
It was tested on the *6G10A v2* model.

‎labgrid/driver/power/tapo.py

+101
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,101 @@
1+
"""Driver for controlling TP-Link Tapo smart plugs and power strips.
2+
3+
This module provides functionality to control TP-Link Tapo smart power devices through
4+
the kasa library. It supports both single socket devices (like P100) and multi-socket
5+
power strips (like P300).
6+
7+
Features:
8+
- Environment-based authentication using KASA_USERNAME and KASA_PASSWORD
9+
- Support for both single and multi-socket devices
10+
11+
Requirements:
12+
- Valid TP-Link cloud credentials (username/password)
13+
"""
14+
15+
import asyncio
16+
import os
17+
import sys
18+
19+
from kasa import Credentials, Device, DeviceConfig, DeviceConnectionParameters, DeviceEncryptionType, DeviceFamily
20+
21+
22+
def _get_credentials() -> Credentials:
23+
username = os.environ.get("KASA_USERNAME")
24+
password = os.environ.get("KASA_PASSWORD")
25+
if username is None or password is None:
26+
raise EnvironmentError("KASA_USERNAME or KASA_PASSWORD environment variable not set")
27+
return Credentials(username=username, password=password)
28+
29+
30+
def _get_connection_type() -> DeviceConnectionParameters:
31+
# Somewhere between python-kasa 0.7.7 and 0.10.2 the API changed
32+
# Labgrid on Python <= 3.10 uses python-kasa 0.7.7
33+
# Labgrid on Python >= 3.11 uses python-kasa 0.10.2
34+
if sys.version_info < (3, 11):
35+
return DeviceConnectionParameters(
36+
device_family=DeviceFamily.SmartTapoPlug,
37+
encryption_type=DeviceEncryptionType.Klap,
38+
https=False,
39+
login_version=2,
40+
)
41+
return DeviceConnectionParameters(
42+
device_family=DeviceFamily.SmartTapoPlug,
43+
encryption_type=DeviceEncryptionType.Klap,
44+
https=False,
45+
login_version=2,
46+
http_port=80,
47+
)
48+
49+
50+
def _get_device_config(host: str) -> DeviceConfig:
51+
# Same as with `_get_connection_type` - python-kasa API changed
52+
if sys.version_info < (3, 11):
53+
return DeviceConfig(
54+
host=host, credentials=_get_credentials(), connection_type=_get_connection_type(), uses_http=True
55+
)
56+
return DeviceConfig(
57+
host=host, credentials=_get_credentials(), connection_type=_get_connection_type()
58+
)
59+
60+
61+
async def _power_set(host: str, port: str, index: str, value: bool) -> None:
62+
"""We embed the coroutines in an `async` function to minimise calls to `asyncio.run`"""
63+
assert port is None
64+
index = int(index)
65+
device = await Device.connect(config=_get_device_config(host))
66+
await device.update()
67+
68+
if device.children:
69+
assert len(device.children) > index, "Trying to access non-existant plug socket on device"
70+
71+
target = device if not device.children else device.children[index]
72+
if value:
73+
await target.turn_on()
74+
else:
75+
await target.turn_off()
76+
await device.disconnect()
77+
78+
79+
def power_set(host: str, port: str, index: str, value: bool) -> None:
80+
asyncio.run(_power_set(host, port, index, value))
81+
82+
83+
async def _power_get(host: str, port: str, index: str) -> bool:
84+
assert port is None
85+
index = int(index)
86+
device = await Device.connect(config=_get_device_config(host))
87+
await device.update()
88+
89+
pwr_state: bool
90+
# If the device has no children, it is a single plug socket
91+
if not device.children:
92+
pwr_state = device.is_on
93+
else:
94+
assert len(device.children) > index, "Trying to access non-existant plug socket on device"
95+
pwr_state = device.children[index].is_on
96+
await device.disconnect()
97+
return pwr_state
98+
99+
100+
def power_get(host: str, port: str, index: str) -> bool:
101+
return asyncio.run(_power_get(host, port, index))

‎pyproject.toml

+1
Original file line numberDiff line numberDiff line change
@@ -81,6 +81,7 @@ dev = [
8181

8282
# additional dev dependencies
8383
"psutil>=5.8.0",
84+
"pytest-asyncio>=0.25.3",
8485
"pytest-benchmark>=4.0.0",
8586
"pytest-cov>=3.0.0",
8687
"pytest-dependency>=0.5.1",

‎tests/test_powerdriver.py

+4
Original file line numberDiff line numberDiff line change
@@ -299,6 +299,10 @@ def test_import_backend_tplink(self):
299299
pytest.importorskip("kasa")
300300
import labgrid.driver.power.tplink
301301

302+
def test_import_backend_tapo(self):
303+
pytest.importorskip("kasa")
304+
import labgrid.driver.power.tapo
305+
302306
def test_import_backend_siglent(self):
303307
pytest.importorskip("vxi11")
304308
import labgrid.driver.power.siglent

‎tests/test_tapo.py

+142
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,142 @@
1+
import os
2+
from unittest.mock import AsyncMock, patch
3+
4+
import pytest
5+
6+
from labgrid.driver.power.tapo import _get_credentials, _power_get, _power_set, power_get
7+
8+
9+
@pytest.fixture
10+
def mock_device_strip():
11+
device = AsyncMock()
12+
device.children = [
13+
AsyncMock(is_on=True),
14+
AsyncMock(is_on=False),
15+
AsyncMock(is_on=True)
16+
]
17+
return device
18+
19+
20+
@pytest.fixture
21+
def mock_device_single_plug():
22+
device = AsyncMock()
23+
device.children = []
24+
return device
25+
26+
27+
@pytest.fixture
28+
def mock_env():
29+
os.environ['KASA_USERNAME'] = 'test_user'
30+
os.environ['KASA_PASSWORD'] = 'test_pass'
31+
yield
32+
del os.environ['KASA_USERNAME']
33+
del os.environ['KASA_PASSWORD']
34+
35+
36+
class TestTapoPowerDriver:
37+
def test_get_credentials_should_raise_value_error_when_credentials_missing(self):
38+
# Save existing environment variables
39+
saved_username = os.environ.pop('KASA_USERNAME', None)
40+
saved_password = os.environ.pop('KASA_PASSWORD', None)
41+
42+
try:
43+
with pytest.raises(EnvironmentError, match="KASA_USERNAME or KASA_PASSWORD environment variable not set"):
44+
_get_credentials()
45+
finally:
46+
# Restore environment variables if they existed
47+
if saved_username is not None:
48+
os.environ['KASA_USERNAME'] = saved_username
49+
if saved_password is not None:
50+
os.environ['KASA_PASSWORD'] = saved_password
51+
52+
def test_credentials_valid(self, mock_env):
53+
creds = _get_credentials()
54+
assert creds.username == 'test_user'
55+
assert creds.password == 'test_pass'
56+
57+
@pytest.mark.asyncio
58+
async def test_power_get_single_plug_turn_on(self, mock_device_single_plug, mock_env):
59+
mock_device_single_plug.is_on = True
60+
61+
with patch('kasa.Device.connect', return_value=mock_device_single_plug):
62+
result = await _power_get('192.168.1.100', None, "0")
63+
assert result is True
64+
65+
@pytest.mark.asyncio
66+
async def test_power_get_single_plug_turn_off(self, mock_device_single_plug, mock_env):
67+
mock_device_single_plug.is_on = False
68+
69+
with patch('kasa.Device.connect', return_value=mock_device_single_plug):
70+
result = await _power_get('192.168.1.100', None, "0")
71+
assert result is False
72+
73+
@pytest.mark.asyncio
74+
async def test_power_get_single_plug_should_not_care_for_index(self, mock_device_single_plug, mock_env):
75+
invalid_index_ignored = "7"
76+
mock_device_single_plug.is_on = True
77+
78+
with patch('kasa.Device.connect', return_value=mock_device_single_plug):
79+
result = await _power_get('192.168.1.100', None, invalid_index_ignored)
80+
assert result is True
81+
82+
@pytest.mark.asyncio
83+
async def test_power_set_single_plug_turn_on(self, mock_device_single_plug, mock_env):
84+
mock_device_single_plug.is_on = False
85+
with patch('kasa.Device.connect', return_value=mock_device_single_plug):
86+
await _power_set('192.168.1.100', None, "0", True)
87+
mock_device_single_plug.turn_on.assert_called_once()
88+
89+
@pytest.mark.asyncio
90+
async def test_power_set_single_plug_turn_off(self, mock_device_single_plug, mock_env):
91+
mock_device_single_plug.is_on = True
92+
with patch('kasa.Device.connect', return_value=mock_device_single_plug):
93+
await _power_set('192.168.1.100', None, "0", False)
94+
mock_device_single_plug.turn_off.assert_called_once()
95+
96+
@pytest.mark.asyncio
97+
async def test_power_get_strip_valid_socket(self, mock_device_strip, mock_env):
98+
with patch('kasa.Device.connect', return_value=mock_device_strip):
99+
# Test first outlet (on)
100+
result = await _power_get('192.168.1.100', None, "0")
101+
assert result is True
102+
103+
# Test second outlet (off)
104+
result = await _power_get('192.168.1.100', None, "1")
105+
assert result is False
106+
107+
# Test third outlet (on)
108+
result = await _power_get('192.168.1.100', None, "2")
109+
assert result is True
110+
111+
@pytest.mark.asyncio
112+
async def test_power_set_strip_valid_socket(self, mock_device_strip, mock_env):
113+
with patch('kasa.Device.connect', return_value=mock_device_strip):
114+
await _power_set('192.168.1.100', None, "0", False)
115+
mock_device_strip.children[0].turn_off.assert_called_once()
116+
117+
await _power_set('192.168.1.100', None, "1", True)
118+
mock_device_strip.children[1].turn_on.assert_called_once()
119+
120+
def test_power_get_should_raise_assertion_error_when_invalid_index_strip(self, mock_device_strip, mock_env):
121+
invalid_socket = "5"
122+
with patch('kasa.Device.connect', return_value=mock_device_strip):
123+
with pytest.raises(AssertionError, match="Trying to access non-existant plug socket"):
124+
power_get('192.168.1.100', None, invalid_socket)
125+
126+
@pytest.mark.asyncio
127+
async def test_power_set_should_raise_assertion_error_when_invalid_index_strip(self, mock_device_strip, mock_env):
128+
invalid_socket = "5"
129+
with patch('kasa.Device.connect', return_value=mock_device_strip):
130+
with pytest.raises(AssertionError, match="Trying to access non-existant plug socket"):
131+
await _power_set('192.168.1.100', None, invalid_socket, True)
132+
133+
def test_port_not_none_strip(self, mock_device_strip):
134+
with patch('kasa.Device.connect', return_value=mock_device_strip):
135+
with pytest.raises(AssertionError):
136+
power_get('192.168.1.100', '8080', "0")
137+
138+
def test_port_not_none_single_socket(self, mock_device_single_plug):
139+
mock_device_single_plug.is_on = True
140+
with patch('kasa.Device.connect', return_value=mock_device_single_plug):
141+
with pytest.raises(AssertionError):
142+
power_get('192.168.1.100', '8080', "0")

0 commit comments

Comments
 (0)