test_ha_mqtt_subscribing_performance.py

"""The tests for MQTT subscribing performance.

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

BROKER = "192.168.xxx.xxx" # Change to the IP of your MQTT broker

TEST_CONFIG_CONTACT = {
    "availability": [
        {
            "topic": "zigbee2mqtt/bridge/state",
            "value_template": "{{ value_json.state }}",
        }
    ],
    "device": {
        "identifiers": ["zigbee2mqtt_0xb0ce{}"],
        "manufacturer": "Sengled",
        "model": "Smart window and door sensor (E1D-G73WNA)",
        "name": "Office Closet Door",
    },
    "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_BATTERY = {
    "availability": [
        {
            "topic": "zigbee2mqtt/bridge/state",
            "value_template": "{{ value_json.state }}",
        }
    ],
    "device": {
        "identifiers": ["zigbee2mqtt_0xb0ce{}"],
        "manufacturer": "Sengled",
        "model": "Smart window and door sensor (E1D-G73WNA)",
        "name": "Door",
    },
    "name": "Door battery {}",
    "state_topic": "zigbee2mqtt/door_{}",
    "unique_id": "0xb0ce{}_battery_zigbee2mqtt",
    "value_template": "{{ value_json.battery }}",
}

CONFIG_TOPIC = "homeassistant/{}/0xb0ce{}/contact/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

COUNT = [1, 512, 1024, 2048, 4096]

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,
) -> None:
    """Test performance subscription."""
    async_fire_mqtt_message(hass, AVAILABILITY_TOPIC, AVAILABILITY_PAYLOAD_OFFLINE)
    sleep(SLEEP1)
    for uniq_id in range(0, count):
        async_fire_mqtt_message(hass, STATE_TOPIC.format(uniq_id), "")
        sleep(0.001)
    sleep(SLEEP2)
    for uniq_id in range(0, count):
        async_fire_mqtt_message(hass, CONFIG_TOPIC.format("binary_sensor", uniq_id), "")
        sleep(0.001)
        async_fire_mqtt_message(hass, CONFIG_TOPIC.format("sensor", uniq_id), "")
        sleep(0.001)


def subscribing_performance(
    count: int,
) -> None:
    """Test performance subscription."""

    def format_config(config: dict[str, Any], config_id: str) -> None:
        """Format the configuration."""
        config["device"]["identifiers"][0] = config["device"]["identifiers"][0].format(
            config_id
        )
        config["unique_id"] = config["unique_id"].format(config_id)
        config["name"] = config["name"].format(config_id)
        config["state_topic"] = config["state_topic"].format(config_id)

    for uniq_id in range(0, count):
        bin_sensor_config = copy.deepcopy(TEST_CONFIG_CONTACT)
        format_config(bin_sensor_config, uniq_id)
        sensor_config = copy.deepcopy(TEST_CONFIG_BATTERY)
        format_config(sensor_config, uniq_id)

        async_fire_mqtt_message(
            hass,
            CONFIG_TOPIC.format("binary_sensor", uniq_id),
            json.dumps(bin_sensor_config),
        )
        sleep(0.001)
        async_fire_mqtt_message(
            hass, CONFIG_TOPIC.format("sensor", uniq_id), json.dumps(sensor_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

    if len(args) == 3 and args[2] == "cleanup":
        clean = True

    if len(args) >= 2 and args[0] == "-n":
        count = int(args[1])
    else:
        print(
            f"Syntax:\n{sys.argv[0]} -n COUNT [cleanup]\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, 1884)
    if clean:
        cleanup(count)
    else:
        subscribing_performance(count)
    client.disconnect()


# Python boilerplate.
if __name__ == "__main__":
    main()