diff --git a/pinecone/grpc/future.py b/pinecone/grpc/future.py index d98e1e84..14fe9cf0 100644 --- a/pinecone/grpc/future.py +++ b/pinecone/grpc/future.py @@ -30,18 +30,17 @@ def _sync_state(self, grpc_future): if self.done(): return - if grpc_future.cancelled(): + if grpc_future.running() and not self.running(): + if not self.set_running_or_notify_cancel(): + grpc_future.cancel() + elif grpc_future.cancelled(): self.cancel() - elif grpc_future.exception(timeout=self._default_timeout): - self.set_exception(grpc_future.exception()) elif grpc_future.done(): try: result = grpc_future.result(timeout=self._default_timeout) self.set_result(result) except Exception as e: self.set_exception(e) - elif grpc_future.running(): - self.set_running_or_notify_cancel() def set_result(self, result): if self._result_transformer: diff --git a/tests/unit_grpc/test_futures.py b/tests/unit_grpc/test_futures.py index 1f8e70c0..43960b00 100644 --- a/tests/unit_grpc/test_futures.py +++ b/tests/unit_grpc/test_futures.py @@ -15,6 +15,8 @@ def mock_grpc_future( grpc_future.exception.return_value = exception grpc_future.running.return_value = running grpc_future.result.return_value = result + if exception: + grpc_future.result.side_effect = exception return grpc_future @@ -309,6 +311,7 @@ def test_result_catch_grpc_exceptions(self, mocker): def test_exception_when_done_maps_grpc_exception(self, mocker): grpc_future = mock_grpc_future(mocker, done=True) grpc_future.exception.return_value = FakeGrpcError(mocker) + grpc_future.result.side_effect = grpc_future.exception.return_value future = PineconeGrpcFuture(grpc_future) @@ -467,6 +470,7 @@ def test_concurrent_futures_wait_first_exception(self, mocker): grpc_future2 = mock_grpc_future(mocker, done=True) grpc_future2.exception.return_value = Exception("Simulated gRPC error") + grpc_future2.result.side_effect = grpc_future2.exception.return_value future2 = PineconeGrpcFuture(grpc_future2) from concurrent.futures import wait, FIRST_EXCEPTION