diff --git a/google/api_core/bidi.py b/google/api_core/bidi.py index 78d98b98..4e800c88 100644 --- a/google/api_core/bidi.py +++ b/google/api_core/bidi.py @@ -306,6 +306,8 @@ def close(self): self._request_queue.put(None) self.call.cancel() self._request_generator = None + self._initial_request = None + self._callbacks = [] # Don't set self.call to None. Keep it around so that send/recv can # raise the error. @@ -717,6 +719,7 @@ def stop(self): _LOGGER.warning("Background thread did not exit.") self._thread = None + self._on_response = None @property def is_active(self): diff --git a/tests/unit/test_bidi.py b/tests/unit/test_bidi.py index c196a682..cca9a214 100644 --- a/tests/unit/test_bidi.py +++ b/tests/unit/test_bidi.py @@ -296,6 +296,9 @@ def test_close(self): # ensure the request queue was signaled to stop. assert bidi_rpc.pending_requests == 1 assert bidi_rpc._request_queue.get() is None + # ensure request and callbacks are cleaned up + assert bidi_rpc._initial_request is None + assert not bidi_rpc._callbacks def test_close_no_rpc(self): bidi_rpc = bidi.BidiRpc(None) @@ -623,6 +626,8 @@ def cancel_side_effect(): assert bidi_rpc.pending_requests == 1 assert bidi_rpc._request_queue.get() is None assert bidi_rpc._finalized + assert bidi_rpc._initial_request is None + assert not bidi_rpc._callbacks def test_reopen_failure_on_rpc_restart(self): error1 = ValueError("1") @@ -777,6 +782,7 @@ def on_response(response): consumer.stop() assert consumer.is_active is False + assert consumer._on_response is None def test_wake_on_error(self): should_continue = threading.Event() @@ -884,6 +890,7 @@ def close_side_effect(): consumer.stop() assert consumer.is_active is False + assert consumer._on_response is None # calling stop twice should not result in an error. consumer.stop()