From 688995ba82e29fff045fecf2c93f7e16cd078cde Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Andrius=20=C5=A0tikonas?= <andrius@stikonas.eu>
Date: Fri, 20 Nov 2020 21:41:06 +0000
Subject: [PATCH] Add get_system call.

---
 neohub.py | 55 ++++++++++++++++++++++++++++++++-----------------------
 test.py   |  7 ++++---
 2 files changed, 36 insertions(+), 26 deletions(-)

diff --git a/neohub.py b/neohub.py
index 20ea0a7..2cbdbbd 100644
--- a/neohub.py
+++ b/neohub.py
@@ -13,37 +13,32 @@ class NeoHub:
 
 
     async def connect(self, host='Neo-Hub', port='4242'):
-        self._reader, self._writer = await asyncio.open_connection(host, port)
-        self._firmware_version = await self._firmware()
+        self._host = host
+        self._port = port
 
 
-    async def _send(self, message, expect_return = True):
+    async def _send(self, message):
+        reader, writer = await asyncio.open_connection(self._host, self._port)
         encoded_message = bytearray(json.dumps(message) + "\0\r", "utf-8")
         self._logger.debug(f"Sending message: {encoded_message}")
-        self._writer.write(encoded_message)
-        await self._writer.drain()
+        writer.write(encoded_message)
+        await writer.drain()
 
-        if expect_return:
-            data = await self._reader.read(4096)
-            json_string = data.decode('utf-8')
-            self._logger.debug(f"Received message: {json_string}")
-            return json.loads(json_string)
-
-
-    def firmware(self):
-        return self._firmware_version
+        data = await reader.read(4096)
+        json_string = data.decode('utf-8')
+        self._logger.debug(f"Received message: {json_string}")
+        writer.close()
+        await writer.wait_closed()
+        print(json_string)
+        return json.loads(json_string)
 
 
     async def firmware(self):
-        return self._firmware_version
-
-
-    async def _firmware(self):
         """
         NeoHub firmware version
         """
-
         message = {"FIRMWARE": 0}
+
         result = await self._send(message)
         firmware_version = int(result['firmware version'])
         return firmware_version
@@ -52,11 +47,25 @@ class NeoHub:
     async def reset(self):
         """
         Reboot neohub
+
+        Returns True if Restart is initiated
         """
         message = {"RESET": 0}
-        
-        if self._firmware_version >= 2027:
-            await self._send(message, expect_return = False)
-            return True
+
+        firmware_version = await self.firmware()
+        result = ""
+        if firmware_version >= 2027:
+            result = await self._send(message)
+            return result['Restarting'] == 1
         else:
             return False
+
+
+    async def get_system(self):
+        """
+        Get system wide variables
+        """
+        message = {"GET_SYSTEM": 0}
+
+        data = await self._send(message)
+        return data
diff --git a/test.py b/test.py
index 05e97da..ee0e060 100755
--- a/test.py
+++ b/test.py
@@ -11,9 +11,10 @@ import logging
 
 async def run():
     hub = neohub.NeoHub()
-    connection = await hub.connect()
-    result = await hub.firmware()
-    print(f"Firmware: {result}")
+    await hub.connect()
+    print(await hub.firmware())
+    result = await hub.get_system()
+    print(result)
 
 
 logging.basicConfig(level=logging.DEBUG)