From 8501aab7c3234ac5f845bfd10087854ec3809cca Mon Sep 17 00:00:00 2001
From: clerie <git@clerie.de>
Date: Sun, 16 Jan 2022 22:29:18 +0100
Subject: [PATCH] Handle message parsing in connection

---
 mitel_ommclient2/connection.py | 8 +++++++-
 mitel_ommclient2/session.py    | 6 ++----
 2 files changed, 9 insertions(+), 5 deletions(-)

diff --git a/mitel_ommclient2/connection.py b/mitel_ommclient2/connection.py
index 428824e..359058e 100644
--- a/mitel_ommclient2/connection.py
+++ b/mitel_ommclient2/connection.py
@@ -6,6 +6,8 @@ import socket
 import ssl
 import threading
 
+from . import messages
+
 class Connection:
     """
         Establishes a connection to the OM Application XML Interface
@@ -46,6 +48,8 @@ class Connection:
             :param message: Message string
         """
 
+        message = messages.construct(message)
+
         self._socket.send(message.encode("utf-8") + b"\0")
 
     def _receive_loop(self):
@@ -74,6 +78,8 @@ class Connection:
                 message, buffer = recv_buffer.split(b"\0", 1)
                 recv_buffer = buffer
 
+                message = message.decode("utf-8")
+                message = messages.parse(message)
                 self._received_messages.put(message)
 
     def recv(self):
@@ -86,7 +92,7 @@ class Connection:
         if self._received_messages.empty():
             return None
 
-        return self._received_messages.get().decode("utf-8")
+        return self._received_messages.get()
 
     def close(self):
         """
diff --git a/mitel_ommclient2/session.py b/mitel_ommclient2/session.py
index 81fa41b..1c69478 100644
--- a/mitel_ommclient2/session.py
+++ b/mitel_ommclient2/session.py
@@ -76,8 +76,6 @@ class Session:
                 'PingResp'
         """
 
-        message = messages.construct(request)
-        self._connection.send(message)
+        self._connection.send(request)
 
-        res = self._wait_for_respose()
-        return messages.parse(res)
+        return self._wait_for_respose()