fix: prevent stream closure on exceptions in S3Writer (#331)

Enhanced __exit__ to log exceptions without closing stream, preventing
masked errors during S3 operations. Added logging and unit tests to
verify exception handling behavior.

Author: jet-tong

s3torchconnector/src/s3torchconnector/s3writer.py

@@ -4,9 +4,12 @@
 import io
 from typing import Union
 import threading
+import logging
 
 from s3torchconnectorclient._mountpoint_s3_client import PutObjectStream
 
+logger = logging.getLogger(__name__)
+
 
 class S3Writer(io.BufferedIOBase):
     """A write-only, file like representation of a single object stored in S3."""
@@ -22,7 +25,16 @@ def __enter__(self):
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
-        self.close()
+        """Close stream on normal exit, log any exceptions that occurred."""
+        if exc_type is not None:
+            try:
+                logger.info(
+                    f"Exception occurred before closing stream: {exc_type.__name__}: {exc_val}"
+                )
+            except:
+                pass
+        else:
+            self.close()
 
     def write(
         self,

s3torchconnector/tst/unit/test_s3writer.py

@@ -87,3 +87,25 @@ def test_concurrent_close_calls():
 
     MOCK_STREAM.close.assert_called_once()
     assert writer._closed
+
+
+def test_exit_without_exception():
+    """Test __exit__ method when no exception occurs."""
+    MOCK_STREAM.reset_mock()
+
+    writer = S3Writer(MOCK_STREAM)
+    writer.__exit__(None, None, None)
+
+    MOCK_STREAM.close.assert_called_once()
+
+
+def test_exit_with_exception(caplog):
+    """Test __exit__ method when an exception occurs."""
+    MOCK_STREAM.reset_mock()
+
+    writer = S3Writer(MOCK_STREAM)
+    test_exception = ValueError("Test exception")
+    writer.__exit__(ValueError, test_exception, None)
+
+    # Stream should not be closed on exception
+    MOCK_STREAM.close.assert_not_called()