johnpaulett#44 - stop waiting for recv() when MLLP end bytes already …

…arrived

hl7/client.py

@@ -4,10 +4,11 @@
 import socket
 import sys
 import time
+import typing
 from argparse import ArgumentParser
 
 import hl7
-from hl7.exceptions import CLIException
+from hl7.exceptions import CLIException, MLLPException
 
 SB = b"\x0b"  # <SB>, vertical tab
 EB = b"\x1c"  # <EB>, file separator
@@ -20,10 +21,6 @@
 log = logging.getLogger(__name__)
 
 
-class MLLPException(Exception):
-    pass
-
-
 class MLLPClient(object):
     """
     A basic, blocking, HL7 MLLP client based upon :py:mod:`socket`.
@@ -68,7 +65,7 @@ def close(self):
         """Release the socket connection"""
         self.socket.close()
 
-    def send_message(self, message):
+    def send_message(self, message: typing.Union[bytes, str, hl7.Message]) -> bytes:
         """Wraps a byte string, unicode string, or :py:class:`hl7.Message`
         in a MLLP container and send the message to the server
 
@@ -90,7 +87,7 @@ def send_message(self, message):
         data = SB + binary + EB + CR
         return self.send(data)
 
-    def send(self, data):
+    def send(self, data: bytes) -> bytes:
         """Low-level, direct access to the socket.send (data must be already
         wrapped in an MLLP container).  Blocks until the server returns.
         """
@@ -112,11 +109,19 @@ def send(self, data):
                 data = self.socket.recv(RECV_BUFFER)
             except TimeoutError:
                 data = None
-            if not data:
-                continue
-            buff += data
+            if data is not None:
+                buff += data
+            # received LLP end markers
+            if buff.endswith(EB + CR):
+                break
         log.debug(f"received {(buff,)}")
-        return buff
+        return self.clean(buff)
+
+    def clean(self, data: bytes) -> bytes:
+        """Removes LLP bytes from data"""
+        data = data.lstrip(SB)
+        data = data.rstrip(EB + CR)
+        return data
 
 
 # wrappers to make testing easier
@@ -234,7 +239,7 @@ def mllp_send(in_args=None):
         "--quiet",
         action="store_false",
         dest="verbose",
-            default=True,
+        default=True,
         help="do not print status messages to stdout",
     )
     parser.add_argument(
@@ -269,11 +274,16 @@ def mllp_send(in_args=None):
     args = parser.parse_args(in_args[1:])
     if args.version:
         import hl7
+
         stdout(hl7.__version__)
         return
 
     host = args.host[0]
 
+    log.setLevel(logging.INFO)
+    if args.verbose:
+        log.setLevel(logging.DEBUG)
+
     if args.filename is not None:
         # Previously set stream to the open() handle, but then we did not
         # close the open file handle.  This new approach consumes the entire
@@ -287,7 +297,6 @@ def mllp_send(in_args=None):
             raise CLIException(1)
 
         stream = stdin()
-
     with MLLPClient(
         host, args.port, deadline=args.deadline, timeout=args.timeout
     ) as client:

hl7/exceptions.py

@@ -18,6 +18,10 @@ class ParseException(HL7Exception):
     pass
 
 
+class MLLPException(HL7Exception):
+    pass
+
+
 class CLIException(HL7Exception):
     """ An exception to propagate expected exit code from cli script"""
     def __init__(self, exit_code):

tests/test_client.py

@@ -1,20 +1,20 @@
+import logging
 import os
 import socket
 import typing
-import logging
 from shutil import rmtree
 from tempfile import mkdtemp
 from unittest import TestCase
 from unittest.mock import patch
 
 import hl7
 from hl7 import __version__ as hl7_version
-from hl7.client import CR, EB, MLLPClient, MLLPException, SB, mllp_send
-from hl7.exceptions import CLIException
-
+from hl7.client import CR, EB, MLLPClient, SB, mllp_send
+from hl7.exceptions import CLIException, MLLPException
 
 log = logging.getLogger(__name__)
 
+
 def return_values_list(*values) -> typing.Iterable[typing.Any]:
     """
     Generates an iterator, which will return each of values and None after values is depleted.
@@ -39,7 +39,6 @@ def setUp(self):
         # use a mock version of socket
         self.socket_patch = patch("hl7.client.socket.socket")
         self.mock_socket = self.socket_patch.start()
-
         self.client = MLLPClient("localhost", 6666, deadline=0.0001)
 
     def tearDown(self):
@@ -57,7 +56,7 @@ def test_close(self):
     def test_send(self):
         # socket.recv returns bytes: https://docs.python.org/3/library/socket.html#socket.socket.recv
         # > Receive data from the socket. The return value is a bytes object representing the data received.
-        self.client.socket.recv.side_effect = return_values_list(b"thanks")
+        self.client.socket.recv.side_effect = return_values_list(b"thanks", EB, CR)
 
         result = self.client.send("foobar\n")
         self.assertEqual(result, b"thanks")
@@ -66,23 +65,23 @@ def test_send(self):
         self.client.socket.recv.assert_any_call(4096)
 
     def test_send_message_unicode(self):
-        self.client.socket.recv.side_effect = return_values_list(b"thanks")
+        self.client.socket.recv.side_effect = return_values_list(b"thanks", EB, CR)
 
         result = self.client.send_message("foobar")
         self.assertEqual(result, b"thanks")
 
         self.client.socket.send.assert_called_once_with(b"\x0bfoobar\x1c\x0d")
 
     def test_send_message_bytestring(self):
-        self.client.socket.recv.side_effect = return_values_list(b"thanks")
+        self.client.socket.recv.side_effect = return_values_list(b"thanks", EB, CR)
 
         result = self.client.send_message(b"foobar")
         self.assertEqual(result, b"thanks")
 
         self.client.socket.send.assert_called_once_with(b"\x0bfoobar\x1c\x0d")
 
     def test_send_message_hl7_message(self):
-        self.client.socket.recv.side_effect = return_values_list(b"thanks")
+        self.client.socket.recv.side_effect = return_values_list(b"thanks", EB, CR)
 
         message = hl7.parse(r"MSH|^~\&|GHH LAB|ELAB")
 
@@ -114,7 +113,7 @@ def setUp(self):
         # patch to avoid touching sys and socket
         self.socket_patch = patch("hl7.client.socket.socket")
         self.mock_socket = self.socket_patch.start()
-        self.mock_socket().recv.side_effect = return_values_list(b"thanks")
+        self.mock_socket().recv.side_effect = return_values_list(b"thanks", EB, CR)
 
         self.stdout_patch = patch("hl7.client.stdout")
         self.mock_stdout = self.stdout_patch.start()
@@ -145,7 +144,7 @@ def setUp(self):
         ]
 
     def _mllp_send(self, args: typing.Optional[typing.List] = None):
-        log.debug('calling mllp_send with args: ', args or self.option_values)
+        log.debug("calling mllp_send with args: ", args or self.option_values)
         return mllp_send(args or self.option_values)
 
     def tearDown(self):
@@ -172,7 +171,7 @@ def test_send(self):
         self.assertFalse(self.mock_exit.called)
 
     def test_send_multiple(self):
-        self.mock_socket().recv.side_effect = return_values_list(b"thanks")
+        self.mock_socket().recv.side_effect = return_values_list(b"thanks", EB, CR)
         self.write(SB + b"foobar" + EB + CR + SB + b"hello" + EB + CR)
 
         self._mllp_send()
@@ -193,7 +192,7 @@ def test_leftover_buffer(self):
 
     def test_quiet(self):
         options = self.option_values.copy()
-        options.append('--quiet')
+        options.append("--quiet")
 
         self._mllp_send(options)
 
@@ -202,7 +201,7 @@ def test_quiet(self):
 
     def test_port(self):
         # replace default port with some exotic value
-        options = self.option_values[:4] + ['7890'] + self.option_values[5:]
+        options = self.option_values[:4] + ["7890"] + self.option_values[5:]
 
         self._mllp_send(options)
 
@@ -221,7 +220,7 @@ def test_stdin(self):
 
     def test_loose_no_stdin(self):
         options = self.option_values.copy()
-        options.append('--loose')
+        options.append("--loose")
         # cut out file path
         options = options[:1] + options[3:]
         self.mock_stdin.return_value = FakeStream()
@@ -233,7 +232,7 @@ def test_loose_no_stdin(self):
 
     def test_loose_windows_newline(self):
         options = self.option_values.copy()
-        options.append('--loose')
+        options.append("--loose")
 
         self.write(SB + b"MSH|^~\\&|foo\r\nbar\r\n" + EB + CR)
 
@@ -245,7 +244,7 @@ def test_loose_windows_newline(self):
 
     def test_loose_unix_newline(self):
         options = self.option_values.copy()
-        options.append('--loose')
+        options.append("--loose")
 
         self.write(SB + b"MSH|^~\\&|foo\nbar\n" + EB + CR)
 
@@ -257,7 +256,7 @@ def test_loose_unix_newline(self):
 
     def test_loose_no_mllp_characters(self):
         options = self.option_values.copy()
-        options.append('--loose')
+        options.append("--loose")
         self.write(b"MSH|^~\\&|foo\r\nbar\r\n")
 
         self._mllp_send(options)
@@ -268,8 +267,8 @@ def test_loose_no_mllp_characters(self):
 
     def test_loose_send_mutliple(self):
         options = self.option_values.copy()
-        options.append('--loose')
-        self.mock_socket().recv.side_effect = return_values_list(b"thanks")
+        options.append("--loose")
+        self.mock_socket().recv.side_effect = return_values_list(b"thanks", EB, CR)
         self.write(b"MSH|^~\\&|1\r\nOBX|1\r\nMSH|^~\\&|2\r\nOBX|2\r\n")
 
         self._mllp_send(options)
@@ -283,10 +282,28 @@ def test_loose_send_mutliple(self):
             SB + b"MSH|^~\\&|2\rOBX|2" + EB + CR,
         )
 
-    def test_version(self):
+    def test_client_end_of_message_parsing(self):
+        options = self.option_values.copy()
+        options.append("--loose")
+        # self.mock_socket().recv.side_effect = return_values_list(b"thanks", EB, CR)
+        self.write(b"MSH|^~\\&|1\r\nOBX|1\r\n")
+
+        self._mllp_send(options)
 
-        options = self.option_values
-        options.append('--version')
+        self.assertEqual(
+            self.mock_socket().recv.call_count, 3, self.mock_socket().recv.mock_calls
+        )
+        # 3 calls for recv - one with the ack, two with closing bytes
+        self.assertEqual(
+            len(self.mock_socket().recv.mock_calls),
+            3,
+            self.mock_socket().recv.mock_calls,
+        )
+        self.mock_stdout.assert_called_once_with(b"thanks")
+
+    def test_version(self):
+        options = self.option_values.copy()
+        options.append("--version")
 
         self._mllp_send(options)