From 71d168511055d81927ad7b15f35d55270efbab4a Mon Sep 17 00:00:00 2001
From: Daniel McKnight <daniel@neon.ai>
Date: Tue, 10 Dec 2024 13:25:33 -0800
Subject: [PATCH] Update documentation to describe multi-part response behavior
 Add test for `stream_callback` final response

---
 README.md                               | 23 ++++++++++++++++++++++-
 neon_mq_connector/utils/client_utils.py |  4 +++-
 tests/test_utils.py                     |  3 +++
 3 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/README.md b/README.md
index 50ce009..484a5af 100644
--- a/README.md
+++ b/README.md
@@ -65,7 +65,28 @@ A response may be sent via:
                        properties=pika.BasicProperties(expiration='1000')
                        )
 ```
-Where `<queue>` is the queue to which the response will be published, and `data` is a `bytes` response (generally a `base64`-encoded `dict`).
+Where `<queue>` is the queue to which the response will be published, and `data`
+is a `bytes` response (generally a `base64`-encoded `dict`).
+
+#### Multi-part Responses
+A callback function may choose to publish multiple response messages so the client
+may receive partial responses as they are being generated. If multiple responses
+will be returned, the following requirements must be met:
+- Each response must be a dict with `_part` and `_is_final` keys.
+- The final response must specify `_is_final=True`
+- The final response *MUST NOT* require the client to handle partial responses
+
+## Client Requests
+Most client applications will interact with services via `send_mq_request`. This
+function will return a `dict` response to the input message.
+
+### Multi-part Responses
+A caller may optionally include a `stream_callback` argument which may receive
+partial responses if supported by the service generating the response. The
+`stream_callback` will always be called with the final result that is returned
+by `send_mq_request`. Keep in mind that the `timeout` param passed to 
+`send_mq_request` applies to the full response, so it may be desirable to increase
+the timeout if 
 
 ### [BETA] Asynchronous Consumers
 
diff --git a/neon_mq_connector/utils/client_utils.py b/neon_mq_connector/utils/client_utils.py
index 9cf368c..0a33369 100644
--- a/neon_mq_connector/utils/client_utils.py
+++ b/neon_mq_connector/utils/client_utils.py
@@ -71,7 +71,9 @@ def send_mq_request(vhost: str, request_data: dict, target_queue: str,
     :param target_queue: queue to post request to
     :param response_queue: optional queue to monitor for a response.
         Generally should be blank
-    :param timeout: time in seconds to wait for a response before timing out
+    :param timeout: time in seconds to wait for a complete response before
+        timing out. Note that in the event of a timeout, a partial response may
+        have been handled by `stream_callback`.
     :param expect_response: boolean indicating whether a response is expected
     :param stream_callback: Optional function to pass partial responses to
     :return: response to request
diff --git a/tests/test_utils.py b/tests/test_utils.py
index c234b46..b25d481 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -200,6 +200,9 @@ def test_multi_part_mq_response(self):
         self.assertEqual(len(response['response'].split()), request['num_parts'])
         self.assertTrue(response['_is_final'])
 
+        # Last callback is the same as the standard response
+        self.assertEqual(response, stream_callback.call_args[0][0])
+
     def test_multiple_mq_requests(self):
         from neon_mq_connector.utils.client_utils import send_mq_request
         responses = dict()