Skip to content

Instantly share code, notes, and snippets.

@jbouwh
Last active May 28, 2024 18:22
Show Gist options
  • Select an option

  • Save jbouwh/10c9ad76a634588bbe891b7e6fd1b6bd to your computer and use it in GitHub Desktop.

Select an option

Save jbouwh/10c9ad76a634588bbe891b7e6fd1b6bd to your computer and use it in GitHub Desktop.
"""Script to test the subscribing performance with device based discovery.
See: https://github.com/home-assistant/core/pull/109030
Test script to to assert the HomeAssistant MQTT discovery.
"""
import copy
import json
import sys
from typing import Any
from time import sleep
import paho.mqtt.client as mqtt
# Change the broker IP address to that
BROKER = "192.168.1.x"
PORT = 1883
AVAILABILITY_DEVICE_ORIGIN = {
"origin": {
"name": "Zigbee2MQTT",
"sw_version": "1.0",
},
"availability": [
{
"topic": "zigbee2mqtt/bridge/state",
"value_template": "{{ value_json.state }}",
}
],
"device": {
"identifiers": ["id"],
"manufacturer": "Sengled",
"model": "Smart window and door sensor (E1D-G73WNA)",
"name": "Office Closet Door",
},
}
TEST_CONFIG_CONTACT = {
"device_class": "door",
"name": "Door contact {}",
"payload_off": True,
"payload_on": False,
"state_topic": "zigbee2mqtt/door_{}",
"unique_id": "0xb0ce_contact_zigbee2mqtt_{}",
"value_template": "{{ value_json.contact }}",
}
TEST_CONFIG_CONTACT_SINGLE = TEST_CONFIG_CONTACT | AVAILABILITY_DEVICE_ORIGIN
TEST_CONFIG_BATTERY = {
"name": "Battery {}",
"device_class": "battery",
"state_topic": "zigbee2mqtt/door_{}",
"unique_id": "0xb0ce_battery_zigbee2mqtt_{}",
"value_template": "{{ value_json.battery }}",
}
TEST_CONFIG_BATTERY_SINGLE = TEST_CONFIG_BATTERY | AVAILABILITY_DEVICE_ORIGIN
TEST_CONFIG_BATTERY_LOW = {
"name": "Battery low {}",
"payload_off": False,
"payload_on": True,
"device_class": "battery",
"state_topic": "zigbee2mqtt/door_{}",
"unique_id": "0xb0ce_battery_zigbee2mqtt_{}",
"value_template": "{{ value_json.battery_low }}",
}
TEST_CONFIG_BATTERY_LOW_SINGLE = TEST_CONFIG_BATTERY_LOW | AVAILABILITY_DEVICE_ORIGIN
TEST_CONFIG_LINKQUALITY = {
"name": "Link quality {}",
"state_topic": "zigbee2mqtt/door_{}",
"unique_id": "0xb0ce_linkquality_zigbee2mqtt_{}",
"value_template": "{{ value_json.linkquality }}",
}
TEST_CONFIG_LINKQUALITY_SINGLE = TEST_CONFIG_LINKQUALITY | AVAILABILITY_DEVICE_ORIGIN
TEST_CONFIG_BUTTON = {
"automation_type": "trigger",
"payload": "short_press",
"topic": "zigbee2mqtt/button_{}",
"type": "button_short_press",
"subtype": "button_1",
}
TEST_CONFIG_BUTTON_SINGLE = TEST_CONFIG_BUTTON | AVAILABILITY_DEVICE_ORIGIN
CONFIG_TOPIC = "homeassistant/{}/0xb0ce{}/{}/config"
DEVICE_CONFIG_TOPIC = "homeassistant/{}/0xb0ce{}/config"
AVAILABILITY_TOPIC = "zigbee2mqtt/bridge/state"
AVAILABILITY_PAYLOAD_ONLINE = '{"state":"online"}'
AVAILABILITY_PAYLOAD_OFFLINE = '{"state":"offline"}'
STATE_TOPIC = "zigbee2mqtt/door_{}"
STATE_PAYLOAD = (
'{"battery":40.0,"battery_low":false,"contact":true,"linkquality":140,'
'"tamper":false,"update":{"installed_version":-1,"latest_version":-1,"state":null},'
'"update_available":null,"voltage":null}'
)
SLEEP1 = 1.0
SLEEP2 = 1.0
SLEEP3 = 1.0
hass = object()
client = mqtt.Client()
def async_fire_mqtt_message(_hass, topic, payload):
"""Publish to broker."""
client.publish(topic, payload, qos=0, retain=True)
def cleanup(count: int, mode: str) -> None:
"""Test performance subscription."""
async_fire_mqtt_message(hass, AVAILABILITY_TOPIC, AVAILABILITY_PAYLOAD_OFFLINE)
sleep(SLEEP1)
for uniq_id in range(0, count):
if mode == "combined":
async_fire_mqtt_message(
hass, DEVICE_CONFIG_TOPIC.format("device", uniq_id), ""
)
sleep(0.001)
else:
async_fire_mqtt_message(
hass, CONFIG_TOPIC.format("binary_sensor", uniq_id, "contact"), ""
)
sleep(0.001)
async_fire_mqtt_message(
hass, CONFIG_TOPIC.format("sensor", uniq_id, "battery"), ""
)
sleep(0.001)
async_fire_mqtt_message(
hass, CONFIG_TOPIC.format("binary_sensor", uniq_id, "battery_low"), ""
)
sleep(0.001)
async_fire_mqtt_message(
hass, CONFIG_TOPIC.format("sensor", uniq_id, "linkquality"), ""
)
sleep(0.001)
async_fire_mqtt_message(
hass, CONFIG_TOPIC.format("device_automation", uniq_id, "button"), ""
)
sleep(0.001)
sleep(SLEEP2)
for uniq_id in range(0, count):
async_fire_mqtt_message(hass, STATE_TOPIC.format(uniq_id), "")
sleep(0.001)
def subscribing_performance(count: int, mode: str) -> None:
"""Test performance subscription."""
def format_config(config: dict[str, Any], config_id: str, state_topic_key: str = "state_topic") -> None:
"""Format the configuration."""
config["device"]["identifiers"][0] = f"zigbee2mqtt_0xb0ce_{config_id}"
config["device"]["name"] = f"Office Closet Door {config_id}"
if "name" in config:
config["name"] = config["name"].format(config_id)
if "unique_id" in config:
config["unique_id"] = config["unique_id"].format(config_id)
config[state_topic_key] = config[state_topic_key].format(config_id)
for uniq_id in range(0, count):
bin_sensor_config_1 = copy.deepcopy(TEST_CONFIG_CONTACT_SINGLE)
format_config(bin_sensor_config_1, uniq_id)
bin_sensor_config_2 = copy.deepcopy(TEST_CONFIG_BATTERY_LOW_SINGLE)
format_config(bin_sensor_config_2, uniq_id)
sensor_config_1 = copy.deepcopy(TEST_CONFIG_BATTERY_SINGLE)
format_config(sensor_config_1, uniq_id)
sensor_config_2 = copy.deepcopy(TEST_CONFIG_LINKQUALITY_SINGLE)
format_config(sensor_config_2, uniq_id)
button_config = copy.deepcopy(TEST_CONFIG_BUTTON_SINGLE)
format_config(button_config, uniq_id, "topic")
if mode == "combined":
for drop_key in ("device", "availability", "origin"):
bin_sensor_config_1.pop(drop_key)
bin_sensor_config_2.pop(drop_key)
sensor_config_1.pop(drop_key)
sensor_config_2.pop(drop_key)
button_config.pop(drop_key)
bin_sensor_config_1["platform"] = "binary_sensor"
sensor_config_1["platform"] = "sensor"
bin_sensor_config_2["platform"] = "binary_sensor"
sensor_config_2["platform"] = "sensor"
button_config["platform"] = "device_automation"
combined = AVAILABILITY_DEVICE_ORIGIN | {
"components": {
"contact": bin_sensor_config_1,
"battery": sensor_config_1,
"battery_low": bin_sensor_config_2,
"linkquality": sensor_config_2,
"button": button_config,
}
}
combined["device"]["identifiers"][0] = f"zigbee2mqtt_0xb0ce_{uniq_id}"
combined["device"]["name"] = f"Office Closet Door {uniq_id}"
async_fire_mqtt_message(
hass,
DEVICE_CONFIG_TOPIC.format("device", uniq_id),
json.dumps(combined),
)
else:
async_fire_mqtt_message(
hass,
CONFIG_TOPIC.format("binary_sensor", uniq_id, "contact"),
json.dumps(bin_sensor_config_1),
)
sleep(0.001)
async_fire_mqtt_message(
hass,
CONFIG_TOPIC.format("sensor", uniq_id, "battery"),
json.dumps(sensor_config_1),
)
sleep(0.001)
async_fire_mqtt_message(
hass,
CONFIG_TOPIC.format("binary_sensor", uniq_id, "battery_low"),
json.dumps(bin_sensor_config_2),
)
sleep(0.001)
async_fire_mqtt_message(
hass,
CONFIG_TOPIC.format("sensor", uniq_id, "linkquality"),
json.dumps(sensor_config_2),
)
sleep(0.001)
async_fire_mqtt_message(
hass,
CONFIG_TOPIC.format("device_automation", uniq_id, "button"),
json.dumps(button_config),
)
sleep(0.001)
sleep(SLEEP1)
async_fire_mqtt_message(hass, AVAILABILITY_TOPIC, AVAILABILITY_PAYLOAD_ONLINE)
sleep(SLEEP2)
for uniq_id in range(0, count):
async_fire_mqtt_message(
hass,
STATE_TOPIC.format(uniq_id),
STATE_PAYLOAD,
)
sleep(0.001)
def main():
"""Perform test."""
clean = False
args = sys.argv[1:]
count: int = 0
mode: str = "single"
if len(args) >= 3 and args[2] == "cleanup":
clean = True
print("Cleanup")
if (
len(args) >= 3
and args[2] == "combined"
or len(args) == 4
and args[3] == "combined"
):
print("Combined mode")
mode = "combined"
if len(args) >= 2 and args[0] == "-n":
count = int(args[1])
else:
print(
f"Syntax:\n{sys.argv[0]} -n COUNT [cleanup] [combined]\n- COUNT: "
"The number of retained test doors to add via MQTT discovery.\n\n"
"Adds a contact binary sensor and battery sensor for each device.\n\n"
"The optional `cleanup` suffix cleans up all retained messages\n"
"of a previous test.\n"
"Adds a contact binary sensor and battery sensor\nfor each device"
)
sys.exit(1)
client.connect(BROKER, PORT)
if clean:
cleanup(count, mode)
else:
subscribing_performance(count, mode)
client.disconnect()
# Run test script.
if __name__ == "__main__":
main()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment