Script for home assistant websocket timing analysis

ha-ws-ping.py

#!/usr/bin/env python3
import json
import sys
import time

from datetime import datetime

# python3-websocket on debian, python-websocket-client on Arch
from websocket import (
    create_connection,
)

with open("config.json", "r") as f:
    ha_config = json.load(f)

HA_HOST = ha_config["host"]
HA_TOKEN = ha_config["token"]


FAIL = '\033[91m'
WARNING = '\033[93m'
NORMAL = '\033[m'


class HaClient:
    def __init__(self):
        self.id = 1
        self.connect()

    def connect(self):
        self.ws = create_connection(f"ws://{HA_HOST}/api/websocket")
        if self.receive()["type"] == "auth_required":
            auth = {"type": "auth", "access_token": HA_TOKEN}
            r = self.send(auth)
            if r["type"] != "auth_ok":
                print("Authentication failed")
                sys.exit()

    def send(self, dict_obj):
        self.ws.send(json.dumps(dict_obj))
        return self.receive()

    def receive(self):
        return json.loads(self.ws.recv())

    def ping(self):
        before = datetime.now()
        service_call = {
            "type": "ping",
            "id": self.id,
        }
        print(f"{before} Ping {self.id}")
        r = self.send(service_call)
        if r["id"] == self.id and r['type'] == "pong":
            after = datetime.now()
            took = (after - before).total_seconds()
            abort = False
            color = NORMAL
            if took > 1:
                color = FAIL
                abort = True
            print(
                f"{color}{after} Got pong for {self.id}, took {took * 1000:.2f} ms{NORMAL}"
            )
            if abort:
                print("Exiting")
                sys.exit()
        self.id += 1

    def close(self):
        self.ws.close()


c = HaClient()
while True:
    c.ping()
    time.sleep(.1)
c.close()