Update websocket test script

This commit reworks the websocket test script to be more usable without
editing it.  Namely it:

1. Moves the script to the websockets library, as it seems to be in more
use
2. Implements an argument parser, so the script doesn't require
modification of host and username/password variables to use.
3. Moves to basic auth, which doesn't require a secondary login
4. implements better parsing of the result rather than simply printing
the json scructure directly.

Tested:
Ran ./websocket_test.py --host <myhostname>

Observed streaming sensor values in the form:
System_Airflow       72.00 CFM
Fan_1                6153.00 RPM
System_Airflow       72.22 CFM
Fan_6                6048.00 RPM
System_Airflow       72.24 CFM
Baseboard_12Volt     12.21 Volts
P105_PCH_AUX         1.05 Volts
PSU1_Input_Power     94.62 Watts
PSU2_Output_Current  6.72 Amps

Signed-off-by: Ed Tanous <ed.tanous@intel.com>
Change-Id: I617f73a5b70f61a391e9c61638f0bcf4437bf6ce

1 files changed, 59 insertions, 20 deletions

diff --git a/scripts/websocket_test.py b/scripts/websocket_test.py
index 8d9c8ae687..f99ddac8ea 100644..100755
--- a/scripts/websocket_test.py
+++ b/scripts/websocket_test.py
@@ -1,21 +1,60 @@
-import json
+#!/usr/bin/env python3
+
+# Example of streaming sensor values from openbmc using the /subscribe api
+# requires websockets package to be installed
+
+import asyncio
 import ssl
-import websocket
-
-websocket.enableTrace(True)
-
-ws = websocket.create_connection('wss://10.243.48.93:18080/subscribe',
-                       sslopt={"cert_reqs": ssl.CERT_NONE},
-                       cookie="XSRF-TOKEN=m0KhYNbxFmUEI4Sr1I22; SESSION=0mdwzoQy3gggQxW3vrEw")
-request = json.dumps({
-    "paths": ["/xyz/openbmc_project/logging", "/xyz/openbmc_project/sensors"],
-    "interfaces": ["xyz.openbmc_project.Logging.Entry", "xyz.openbmc_project.Sensor.Value"]
-})
-
-ws.send(request)
-print("Sent")
-print("Receiving...")
-while True:
-    result = ws.recv()
-    print("Received '%s'" % result)
-ws.close()
+import websockets
+import base64
+import json
+import argparse
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--host", help="Host to connect to", required=True)
+parser.add_argument(
+    "--username", help="Username to connect with", default="root")
+parser.add_argument("--password", help="Password to use", default="0penBmc")
+
+args = parser.parse_args()
+
+sensor_type_map = {
+    "voltage": "Volts",
+    "power": "Watts",
+    "fan": "RPM",
+    "fan_tach": "RPM",
+    "temperature": "Degrees C",
+    "altitude": "Meters",
+    "current": "Amps",
+    "energy": "Joules",
+    "cfm": "CFM"
+}
+
+
+async def hello():
+    uri = 'wss://{}/subscribe'.format(args.host)
+    ssl_context = ssl.SSLContext()
+    authbytes = "{}:{}".format(args.username, args.password).encode('ascii')
+    auth = "Basic {}".format(base64.b64encode(authbytes).decode('ascii'))
+    headers = {"Authorization": auth}
+    async with websockets.connect(uri, ssl=ssl_context, extra_headers=headers) as websocket:
+        request = json.dumps({
+            "paths": ["/xyz/openbmc_project/sensors"],
+            "interfaces": ["xyz.openbmc_project.Sensor.Value"]
+        })
+        await websocket.send(request)
+
+        while True:
+            payload = await websocket.recv()
+            j = json.loads(payload)
+            path = j.get("path", "unknown/unknown")
+            name = path.split("/")[-1]
+            sensor_type = path.split("/")[-2]
+            units = sensor_type_map.get(sensor_type, "")
+            properties = j.get("properties", [])
+            value = properties.get("Value", None)
+            if value is None:
+                continue
+            print(f"{name:<20} {value:4.02f} {units}")
+
+asyncio.get_event_loop().run_until_complete(hello())