diff --git a/awscli/customizations/s3/tasks.py b/awscli/customizations/s3/tasks.py index d187fb105530..430004ca3ae8 100644 --- a/awscli/customizations/s3/tasks.py +++ b/awscli/customizations/s3/tasks.py @@ -560,10 +560,10 @@ def wait_for_parts_to_finish(self): def wait_for_upload_id(self): with self._upload_id_condition: - while self._upload_id is None: - if self._state == self._CANCELLED: - raise UploadCancelledError("Upload has been cancelled.") - self._upload_id_condition.wait(timeout=1) + while self._upload_id is None and self._state != self._CANCELLED: + self._upload_id_condition.wait(timeout=1) + if self._state == self._CANCELLED: + raise UploadCancelledError("Upload has been cancelled.") return self._upload_id def wait_for_completion(self): diff --git a/tests/unit/customizations/s3/test_tasks.py b/tests/unit/customizations/s3/test_tasks.py index 051d21035076..4451c85cb569 100644 --- a/tests/unit/customizations/s3/test_tasks.py +++ b/tests/unit/customizations/s3/test_tasks.py @@ -223,6 +223,26 @@ def test_can_cancel_tasks(self): with self.assertRaises(UploadCancelledError): self.context.wait_for_parts_to_finish() + def test_cancel_after_upload_id(self): + # We want have a thread waiting for the upload id. + upload_part_thread = threading.Thread(target=self.upload_part, + args=(1,)) + self.start_thread(upload_part_thread) + + # We announce the upload id. + self.create_upload('my_upload_id') + # The upload_part thread can now proceed, + # now, let's cancel this upload. + self.context.cancel_upload() + + # The upload_part_thread should be finished. + self.join_threads() + + # In a cancelled multipart upload task any subsequent + # call to wait_for_upload_id must raise an UploadCancelledError + with self.assertRaises(UploadCancelledError): + self.context.wait_for_upload_id() + def test_cancel_threads_waiting_for_completion(self): # So we have a thread waiting for the entire upload to complete. arbitrary_waiting_thread = threading.Thread(target=self.wait_for_upload_complete)