Add set_exception for CRT Transfer future

.changes/next-release/enhancement-crt-20073.json

@@ -0,0 +1,5 @@
+{
+  "type": "enhancement",
+  "category": "``crt``",
+  "description": "Add ``set_exception`` to ``CRTTransferFuture`` to allow setting exceptions in subscribers."
+}

s3transfer/crt.py

@@ -27,6 +27,7 @@
 from awscrt.io import ClientTlsContext, TlsContextOptions
 from awscrt.auth import AwsCredentialsProvider, AwsCredentials
 
+from s3transfer.exceptions import TransferNotDoneError
 from s3transfer.futures import BaseTransferFuture, BaseTransferMeta
 from s3transfer.utils import CallArgs, OSUtils, get_callbacks
 from s3transfer.constants import GB, MB
@@ -317,6 +318,14 @@ def result(self, timeout=None):
     def cancel(self):
         self._coordinator.cancel()
 
+    def set_exception(self, exception):
+        """Sets the exception on the future."""
+        if not self.done():
+            raise TransferNotDoneError(
+                'set_exception can only be called once the transfer is '
+                'complete.')
+        self._coordinator.set_exception(exception, override=True)
+
 
 class BaseCRTRequestSerializer:
     def serialize_http_request(self, transfer_type, future):

tests/unit/test_crt.py

@@ -15,15 +15,21 @@
 
 from botocore.session import Session
 from botocore.credentials import CredentialResolver, ReadOnlyCredentials
+from s3transfer.exceptions import TransferNotDoneError
 from s3transfer.utils import CallArgs
 
 from tests import FileCreator
 from tests import requires_crt, HAS_CRT
 
 if HAS_CRT:
+    import awscrt.s3
     import s3transfer.crt
 
 
+class CustomFutureException(Exception):
+    pass
+
+
 @requires_crt
 class TestBotocoreCRTRequestSerializer(unittest.TestCase):
     def setUp(self):
@@ -123,3 +129,31 @@ def test_load_credentials_once(self):
         # will only be called once
         self.assertEqual(
             self.botocore_credential_provider.load_credentials.call_count, 1)
+
+
+@requires_crt
+class TestCRTTransferFuture(unittest.TestCase):
+    def setUp(self):
+        self.mock_s3_request = mock.Mock(awscrt.s3.S3RequestType)
+        self.mock_crt_future = mock.Mock(awscrt.s3.Future)
+        self.mock_s3_request.finished_future = self.mock_crt_future
+        self.coordinator = s3transfer.crt.CRTTransferCoordinator()
+        self.coordinator.set_s3_request(self.mock_s3_request)
+        self.future = s3transfer.crt.CRTTransferFuture(
+            coordinator=self.coordinator)
+
+    def test_set_exception(self):
+        self.future.set_exception(CustomFutureException())
+        with self.assertRaises(CustomFutureException):
+            self.future.result()
+
+    def test_set_exception_raises_error_when_not_done(self):
+        self.mock_crt_future.done.return_value = False
+        with self.assertRaises(TransferNotDoneError):
+            self.future.set_exception(CustomFutureException())
+
+    def test_set_exception_can_override_previous_exception(self):
+        self.future.set_exception(Exception())
+        self.future.set_exception(CustomFutureException())
+        with self.assertRaises(CustomFutureException):
+            self.future.result()