From cb5c7db7f006c4852f3072751c766e629ecb57c3 Mon Sep 17 00:00:00 2001
From: Daniel <d.hornung@indiscale.com>
Date: Mon, 25 Nov 2024 16:23:03 +0100
Subject: [PATCH] FIX: Fixed undefined variable.

---
 src/linkahead/connection/connection.py | 25 ++++++-------
 unittests/test_connection.py           | 52 ++++++++++++++++++++++++++
 2 files changed, 64 insertions(+), 13 deletions(-)

diff --git a/src/linkahead/connection/connection.py b/src/linkahead/connection/connection.py
index f75226d4..4c40842a 100644
--- a/src/linkahead/connection/connection.py
+++ b/src/linkahead/connection/connection.py
@@ -83,8 +83,10 @@ class _WrappedHTTPResponse(CaosDBHTTPResponse):
         return self.response.status_code
 
     def read(self, size: Optional[int] = None):
+        # FIXME This function behaves unexpectedly if `size` is larger than in the first run.
+
         if self._stream_consumed is True:
-            raise RuntimeError("Stream is consumed")
+            raise BufferError("Stream is consumed")
 
         if self._buffer is None:
             # the buffer has been drained in the previous call.
@@ -97,18 +99,15 @@ class _WrappedHTTPResponse(CaosDBHTTPResponse):
             return self.response.content
 
         if size is None or size == 0:
-            raise RuntimeError(
-                "size parameter should not be None if the stream is not consumed yet")
+            raise BufferError(
+                "`size` parameter can not be None or zero once reading has started with a non-zero "
+                "value.")
 
         if len(self._buffer) >= size:
             # still enough bytes in the buffer
-            # FIXME: `chunk`` is used before definition
-            raise NotImplementedError("chunk is undefined")
-
-            # # old code:
-            # result = chunk[:size]
-            # self._buffer = chunk[size:]
-            # return result
+            result = self._buffer[:size]
+            self._buffer = self._buffer[size:]
+            return result
 
         if self._generator is None:
             # first call to this method
@@ -119,16 +118,16 @@ class _WrappedHTTPResponse(CaosDBHTTPResponse):
         try:
             # read new data into the buffer
             chunk = self._buffer + next(self._generator)
-            result = chunk[:size]
+            result = chunk[:size]  # FIXME what if `size` is larger than at `iter_content(size)`?
             if len(result) == 0:
                 self._stream_consumed = True
             self._buffer = chunk[size:]
             return result
         except StopIteration:
             # drain buffer
-            result = self._buffer
+            last_result = self._buffer
             self._buffer = None
-            return result
+            return last_result
 
     def getheader(self, name: str, default=None):
         return self.response.headers[name] if name in self.response.headers else default
diff --git a/unittests/test_connection.py b/unittests/test_connection.py
index a3a1eff7..a0d280c2 100644
--- a/unittests/test_connection.py
+++ b/unittests/test_connection.py
@@ -25,14 +25,18 @@
 # pylint: disable=missing-docstring
 from __future__ import print_function, unicode_literals
 
+import io
 import re
 from builtins import bytes, str  # pylint: disable=redefined-builtin
 
+import requests
+
 from linkahead import execute_query
 from linkahead.configuration import _reset_config, get_config
 from linkahead.connection.authentication.interface import CredentialsAuthenticator
 from linkahead.connection.connection import (CaosDBServerConnection,
                                              _DefaultCaosDBServerConnection,
+                                             _WrappedHTTPResponse,
                                              configure_connection)
 from linkahead.connection.mockup import (MockUpResponse, MockUpServerConnection,
                                          _request_log_message)
@@ -324,3 +328,51 @@ def test_auth_token_connection():
                                 "auth_token authenticator cannot log in "
                                 "again. You must provide a new authentication "
                                 "token.")
+
+
+def test_buffer_read():
+    """Test the buffering in _WrappedHTTPResponse.read()"""
+
+    class MockResponse(requests.Response):
+        def __init__(self, content: bytes):
+            """A mock response
+
+            Parameters
+            ----------
+            content : bytes
+              The fake content.
+            """
+            super().__init__()
+            self._content = content
+            bio = io.BytesIO(expected)
+            self.raw = bio
+
+    expected = b"This response."
+    MockResponse(expected)
+
+    #############################
+    # Check for some exceptions #
+    #############################
+    resp = _WrappedHTTPResponse(response=MockResponse(expected))
+    with raises(BufferError) as rte:
+        resp.read(4)
+        resp.read()
+    assert "`size` parameter can not be None" in str(rte.value)
+
+    resp = _WrappedHTTPResponse(response=MockResponse(expected))
+    with raises(BufferError) as rte:
+        resp.read(4)
+        resp.read(0)
+    assert "`size` parameter can not be None" in str(rte.value)
+
+    print("---")
+    resp = _WrappedHTTPResponse(response=MockResponse(expected))
+    result = (
+        resp.read(4)
+        + resp.read(2)
+        + resp.read(2)  # This line failed before.
+        + resp.read(4)  # Reading the rest in two chunks, because of current limitations in read().
+        + resp.read(2)
+    )
+
+    assert result == expected
-- 
GitLab