chore: add initial tests for s3_seekable_io

Author: rubenfonseca

aws_lambda_powertools/utilities/streaming/_s3_seekable_io.py

@@ -39,6 +39,9 @@ def __init__(
         # Holds the current position in the stream
         self._position = 0
 
+        # Stores the closed state of the stream
+        self._closed: bool = False
+
         # Caches the size of the object
         self._size: Optional[int] = None
 
@@ -88,6 +91,7 @@ def raw_stream(self) -> StreamingBody:
             range_header = "bytes=%d-" % self._position
             logging.debug(f"Starting new stream at {range_header}...")
             self._raw_stream = self.s3_object.get(Range=range_header)["Body"]
+            self._closed = False
 
         return self._raw_stream
 
@@ -140,11 +144,13 @@ def readline(self, size: Optional[int] = None) -> bytes:
 
     def readlines(self, hint: int = -1) -> List[bytes]:
         # boto3's StreamingResponse doesn't implement the "hint" parameter
-        return self.raw_stream.readlines()
+        data = self.raw_stream.readlines()
+        self._position += sum(len(line) for line in data)
+        return data
 
     @property
     def closed(self) -> bool:
-        return self.raw_stream.closed
+        return self._closed
 
     def __next__(self):
         return self.raw_stream.__next__()
@@ -155,11 +161,12 @@ def __iter__(self):
     def __enter__(self):
         return self
 
-    def __exit__(self, **kwargs):
+    def __exit__(self, *kwargs):
         self.close()
 
     def close(self) -> None:
         self.raw_stream.close()
+        self._closed = True
 
     def fileno(self) -> int:
         raise NotImplementedError()

tests/functional/streaming/__init__.py

@@ -0,0 +1,201 @@
+import io
+
+import boto3
+import pytest
+from botocore import stub
+from botocore.response import StreamingBody
+
+from aws_lambda_powertools.utilities.streaming._s3_seekable_io import _S3SeekableIO
+
+
+@pytest.fixture
+def s3_resource():
+    return boto3.resource("s3")
+
+
+@pytest.fixture
+def s3_seekable_obj(s3_resource):
+    return _S3SeekableIO(bucket="bucket", key="key", boto3_s3_resource=s3_resource)
+
+
+@pytest.fixture
+def s3_resource_stub(s3_resource):
+    s3_stub = stub.Stubber(s3_resource.meta.client)
+    s3_stub.activate()
+    return s3_stub
+
+
+def test_seekable(s3_seekable_obj):
+    assert s3_seekable_obj.seekable() is True
+
+
+def test_readable(s3_seekable_obj):
+    assert s3_seekable_obj.readable() is True
+
+
+def test_writeable(s3_seekable_obj):
+    assert s3_seekable_obj.writable() is False
+
+
+def test_tell_is_zero(s3_seekable_obj):
+    assert s3_seekable_obj.tell() == 0
+
+
+def test_seek_set_changes_position(s3_seekable_obj):
+    assert s3_seekable_obj.seek(300, io.SEEK_SET) == 300
+    assert s3_seekable_obj.tell() == 300
+
+
+def test_seek_cur_changes_position(s3_seekable_obj):
+    assert s3_seekable_obj.seek(200, io.SEEK_CUR) == 200
+    assert s3_seekable_obj.seek(100, io.SEEK_CUR) == 300
+    assert s3_seekable_obj.tell() == 300
+
+
+def test_seek_end(s3_seekable_obj, s3_resource_stub):
+    s3_resource_stub.add_response("head_object", {"ContentLength": 1000})
+
+    assert s3_seekable_obj.seek(0, io.SEEK_END) == 1000
+    assert s3_seekable_obj.tell() == 1000
+
+
+def test_size(s3_seekable_obj, s3_resource_stub):
+    s3_resource_stub.add_response("head_object", {"ContentLength": 1000})
+
+    assert s3_seekable_obj.size == 1000
+
+
+def test_raw_stream_fetches_with_range_header(s3_seekable_obj, s3_resource_stub):
+    s3_resource_stub.add_response(
+        "get_object",
+        {"Body": ""},
+        {"Bucket": s3_seekable_obj.bucket, "Key": s3_seekable_obj.key, "Range": "bytes=0-"},
+    )
+
+    assert s3_seekable_obj.raw_stream is not None
+
+
+def test_raw_stream_fetches_with_range_header_after_seek(s3_seekable_obj, s3_resource_stub):
+    s3_seekable_obj.seek(100, io.SEEK_SET)
+
+    s3_resource_stub.add_response(
+        "get_object",
+        {"Body": ""},
+        {"Bucket": s3_seekable_obj.bucket, "Key": s3_seekable_obj.key, "Range": "bytes=100-"},
+    )
+
+    assert s3_seekable_obj.raw_stream is not None
+
+
+def test_read(s3_seekable_obj, s3_resource_stub):
+    payload = b"hello world"
+    streaming_body = StreamingBody(raw_stream=io.BytesIO(payload), content_length=len(payload))
+
+    s3_resource_stub.add_response(
+        "get_object",
+        {"Body": streaming_body},
+        {"Bucket": s3_seekable_obj.bucket, "Key": s3_seekable_obj.key, "Range": "bytes=0-"},
+    )
+
+    assert s3_seekable_obj.read(5) == b"hello"
+    assert s3_seekable_obj.read(1) == b" "
+    assert s3_seekable_obj.read(10) == b"world"
+    assert s3_seekable_obj.tell() == len(payload)
+
+
+def test_readline(s3_seekable_obj, s3_resource_stub):
+    payload = b"hello world\nworld hello"
+    streaming_body = StreamingBody(raw_stream=io.BytesIO(payload), content_length=len(payload))
+
+    s3_resource_stub.add_response(
+        "get_object",
+        {"Body": streaming_body},
+        {"Bucket": s3_seekable_obj.bucket, "Key": s3_seekable_obj.key, "Range": "bytes=0-"},
+    )
+
+    assert s3_seekable_obj.readline() == b"hello world\n"
+    assert s3_seekable_obj.readline() == b"world hello"
+    assert s3_seekable_obj.tell() == len(payload)
+
+
+def test_readlines(s3_seekable_obj, s3_resource_stub):
+    payload = b"hello world\nworld hello"
+    streaming_body = StreamingBody(raw_stream=io.BytesIO(payload), content_length=len(payload))
+
+    s3_resource_stub.add_response(
+        "get_object",
+        {"Body": streaming_body},
+        {"Bucket": s3_seekable_obj.bucket, "Key": s3_seekable_obj.key, "Range": "bytes=0-"},
+    )
+
+    assert s3_seekable_obj.readlines() == [b"hello world\n", b"world hello"]
+    assert s3_seekable_obj.tell() == len(payload)
+
+
+def test_closed(s3_seekable_obj, s3_resource_stub):
+    payload = b"test"
+    streaming_body = StreamingBody(raw_stream=io.BytesIO(payload), content_length=len(payload))
+
+    s3_resource_stub.add_response(
+        "get_object",
+        {"Body": streaming_body},
+        {"Bucket": s3_seekable_obj.bucket, "Key": s3_seekable_obj.key, "Range": "bytes=0-"},
+    )
+
+    s3_seekable_obj.close()
+    assert s3_seekable_obj.closed is True
+
+
+def test_next(s3_seekable_obj, s3_resource_stub):
+    payload = b"test"
+    streaming_body = StreamingBody(raw_stream=io.BytesIO(payload), content_length=len(payload))
+
+    s3_resource_stub.add_response(
+        "get_object",
+        {"Body": streaming_body},
+        {"Bucket": s3_seekable_obj.bucket, "Key": s3_seekable_obj.key, "Range": "bytes=0-"},
+    )
+
+    assert next(s3_seekable_obj) == b"test"
+    with pytest.raises(StopIteration):
+        next(s3_seekable_obj)
+
+
+def test_context_manager(s3_seekable_obj, s3_resource_stub):
+    payload = b"test"
+    streaming_body = StreamingBody(raw_stream=io.BytesIO(payload), content_length=len(payload))
+
+    s3_resource_stub.add_response(
+        "get_object",
+        {"Body": streaming_body},
+        {"Bucket": s3_seekable_obj.bucket, "Key": s3_seekable_obj.key, "Range": "bytes=0-"},
+    )
+
+    with s3_seekable_obj as f:
+        assert f.read(4) == b"test"
+
+    assert s3_seekable_obj.closed is True
+
+
+def test_fileno(s3_seekable_obj):
+    with pytest.raises(NotImplementedError):
+        s3_seekable_obj.fileno()
+
+
+def test_flush(s3_seekable_obj):
+    with pytest.raises(NotImplementedError):
+        s3_seekable_obj.flush()
+
+
+def test_isatty(s3_seekable_obj):
+    assert s3_seekable_obj.isatty() is False
+
+
+def test_truncate(s3_seekable_obj):
+    with pytest.raises(NotImplementedError):
+        s3_seekable_obj.truncate()
+
+
+def test_write(s3_seekable_obj):
+    with pytest.raises(NotImplementedError):
+        s3_seekable_obj.write(b"data")