diff --git a/packages/google-api-core/google/api_core/exceptions.py b/packages/google-api-core/google/api_core/exceptions.py index aa9898c5b9a6..120736e2169e 100644 --- a/packages/google-api-core/google/api_core/exceptions.py +++ b/packages/google-api-core/google/api_core/exceptions.py @@ -93,6 +93,9 @@ def cause(self): def __str__(self): return "{}, last exception: {}".format(self.message, self.cause) + def __reduce__(self): + return (type(self), (self.message, self.cause), self.__dict__) + class _GoogleAPICallErrorMeta(type): """Metaclass for registering GoogleAPICallError subclasses.""" diff --git a/packages/google-api-core/tests/unit/test_exceptions.py b/packages/google-api-core/tests/unit/test_exceptions.py index 6a71ddeb9261..31b592639e33 100644 --- a/packages/google-api-core/tests/unit/test_exceptions.py +++ b/packages/google-api-core/tests/unit/test_exceptions.py @@ -12,8 +12,10 @@ # See the License for the specific language governing permissions and # limitations under the License. +import copy import http.client import json +import pickle from unittest import mock import pytest @@ -434,3 +436,42 @@ def test_from_grpc_error_misleading_404(): assert exception.grpc_status_code == grpc.StatusCode.UNIMPLEMENTED assert "Received http2 header with status: 404" in exception.message assert "api_endpoint" in exception.message + + +@pytest.mark.parametrize( + "operation", + [ + lambda err: pickle.loads(pickle.dumps(err)), + copy.copy, + copy.deepcopy, + ], + ids=["pickle", "copy", "deepcopy"], +) +@pytest.mark.parametrize( + "cause", + [ + exceptions.ServiceUnavailable("503 backend unavailable"), + None, + ], + ids=["with_cause", "none_cause"], +) +def test_retry_error_copy_and_pickle(operation, cause): + original = exceptions.RetryError("Request timed out", cause) + original.custom_attr = "test_value" + + result = operation(original) + + assert isinstance(result, exceptions.RetryError) + assert result.message == original.message + assert str(result) == str(original) + assert result.custom_attr == "test_value" + + if cause is None: + assert result.cause is None + else: + if operation is copy.copy: + assert result.cause is original.cause + + assert isinstance(result.cause, type(original.cause)) + assert result.cause.message == original.cause.message + assert str(result.cause) == str(original.cause)