From 9ba3f624d1fc8cd543ef4cdad013f37976b436dc Mon Sep 17 00:00:00 2001
From: Vincent Bernat <vincent@bernat.ch>
Date: Thu, 31 Mar 2022 09:21:54 +0200
Subject: [PATCH] i3-companion: reduce a bit the code to query battery status

---
 bin/i3-companion | 21 ++++++++++-----------
 1 file changed, 10 insertions(+), 11 deletions(-)

diff --git a/bin/i3-companion b/bin/i3-companion
index adf2278..a643688 100755
--- a/bin/i3-companion
+++ b/bin/i3-companion
@@ -844,17 +844,16 @@ async def bluetooth_status(i3, event, *args):
                     fut = loop.create_future()
                     loop._sock_connect(fut, sock, (mac, 8))
                     await fut
-                    # Init
-                    await loop.sock_sendall(sock, b"\0\1\1\0")
-                    ack = await loop.sock_recv(sock, 4)
-                    assert ack == b"\0\1\3\5"
-                    await loop.sock_recv(sock, 5)  # firmware
-                    # Battery level
-                    await loop.sock_sendall(sock, b"\2\2\1\0")
-                    ack = await loop.sock_recv(sock, 4)
-                    assert ack == b"\2\2\3\1"
-                    battery = await loop.sock_recv(sock, 1)
-                    battery = battery[0]
+                    # Get battery
+                    for send, ack, size in (
+                        (b"\0\1\1\0", b"\0\1\3\5", 5),  # get firmware
+                        (b"\2\2\1\0", b"\2\2\3\1", 1),  # get battery
+                    ):
+                        await loop.sock_sendall(sock, send)
+                        rack = await loop.sock_recv(sock, len(ack))
+                        assert rack == ack, "incorrect ack received"
+                        result = await loop.sock_recv(sock, size)
+                    battery = result[0]
                 finally:
                     sock.close()
                 # Choose an icon