From b7c1df52e621b9f4783dcd91a706a5eb9b918a8f Mon Sep 17 00:00:00 2001
From: kyleknap <kyknapp1@gmail.com>
Date: Tue, 20 Apr 2021 08:22:27 -0700
Subject: [PATCH] Add set_exception for CRT Transfer future

This enables more parity with the threaded, Python transfer
manager where a subscriber can set an exception in the on_done
callback such as failing to delete the source file after
transfering it.
---
 .../next-release/enhancement-crt-20073.json   |  5 +++
 s3transfer/crt.py                             |  9 +++++
 tests/unit/test_crt.py                        | 34 +++++++++++++++++++
 3 files changed, 48 insertions(+)
 create mode 100644 .changes/next-release/enhancement-crt-20073.json

diff --git a/.changes/next-release/enhancement-crt-20073.json b/.changes/next-release/enhancement-crt-20073.json
new file mode 100644
index 00000000..c9033dee
--- /dev/null
+++ b/.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."
+}
diff --git a/s3transfer/crt.py b/s3transfer/crt.py
index 7d66f97e..cf68c603 100644
--- a/s3transfer/crt.py
+++ b/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):
diff --git a/tests/unit/test_crt.py b/tests/unit/test_crt.py
index 05adf34a..54e3327c 100644
--- a/tests/unit/test_crt.py
+++ b/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()