Skip to content
This repository has been archived by the owner on Oct 11, 2023. It is now read-only.

feat: add context manager support in client #38

Merged
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
6 changes: 6 additions & 0 deletions google/cloud/debugger_v2/services/controller2/async_client.py
Expand Up @@ -431,6 +431,12 @@ async def update_active_breakpoint(
# Done; return the response.
return response

async def __aenter__(self):
return self

async def __aexit__(self, exc_type, exc, tb):
await self.transport.close()


try:
DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo(
Expand Down
18 changes: 14 additions & 4 deletions google/cloud/debugger_v2/services/controller2/client.py
Expand Up @@ -348,10 +348,7 @@ def __init__(
client_cert_source_for_mtls=client_cert_source_func,
quota_project_id=client_options.quota_project_id,
client_info=client_info,
always_use_jwt_access=(
Transport == type(self).get_transport_class("grpc")
or Transport == type(self).get_transport_class("grpc_asyncio")
),
always_use_jwt_access=True,
)

def register_debuggee(
Expand Down Expand Up @@ -588,6 +585,19 @@ def update_active_breakpoint(
# Done; return the response.
return response

def __enter__(self):
return self

def __exit__(self, type, value, traceback):
"""Releases underlying transport's resources.

.. warning::
ONLY use as a context manager if the transport is NOT shared
with other clients! Exiting the with block will CLOSE the transport
and may cause errors in other clients!
"""
self.transport.close()


try:
DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo(
Expand Down
Expand Up @@ -192,6 +192,15 @@ def _prep_wrapped_messages(self, client_info):
),
}

def close(self):
"""Closes resources associated with the transport.

.. warning::
Only call this method if the transport is NOT shared
with other clients - this may cause errors in other clients!
"""
raise NotImplementedError()

@property
def register_debuggee(
self,
Expand Down
Expand Up @@ -364,5 +364,8 @@ def update_active_breakpoint(
)
return self._stubs["update_active_breakpoint"]

def close(self):
self.grpc_channel.close()


__all__ = ("Controller2GrpcTransport",)
Expand Up @@ -368,5 +368,8 @@ def update_active_breakpoint(
)
return self._stubs["update_active_breakpoint"]

def close(self):
return self.grpc_channel.close()


__all__ = ("Controller2GrpcAsyncIOTransport",)
6 changes: 6 additions & 0 deletions google/cloud/debugger_v2/services/debugger2/async_client.py
Expand Up @@ -599,6 +599,12 @@ async def list_debuggees(
# Done; return the response.
return response

async def __aenter__(self):
return self

async def __aexit__(self, exc_type, exc, tb):
await self.transport.close()


try:
DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo(
Expand Down
18 changes: 14 additions & 4 deletions google/cloud/debugger_v2/services/debugger2/client.py
Expand Up @@ -339,10 +339,7 @@ def __init__(
client_cert_source_for_mtls=client_cert_source_func,
quota_project_id=client_options.quota_project_id,
client_info=client_info,
always_use_jwt_access=(
Transport == type(self).get_transport_class("grpc")
or Transport == type(self).get_transport_class("grpc_asyncio")
),
always_use_jwt_access=True,
)

def set_breakpoint(
Expand Down Expand Up @@ -738,6 +735,19 @@ def list_debuggees(
# Done; return the response.
return response

def __enter__(self):
return self

def __exit__(self, type, value, traceback):
"""Releases underlying transport's resources.

.. warning::
ONLY use as a context manager if the transport is NOT shared
with other clients! Exiting the with block will CLOSE the transport
and may cause errors in other clients!
"""
self.transport.close()


try:
DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo(
Expand Down
Expand Up @@ -223,6 +223,15 @@ def _prep_wrapped_messages(self, client_info):
),
}

def close(self):
"""Closes resources associated with the transport.

.. warning::
Only call this method if the transport is NOT shared
with other clients - this may cause errors in other clients!
"""
raise NotImplementedError()

@property
def set_breakpoint(
self,
Expand Down
Expand Up @@ -368,5 +368,8 @@ def list_debuggees(
)
return self._stubs["list_debuggees"]

def close(self):
self.grpc_channel.close()


__all__ = ("Debugger2GrpcTransport",)
Expand Up @@ -379,5 +379,8 @@ def list_debuggees(
)
return self._stubs["list_debuggees"]

def close(self):
return self.grpc_channel.close()


__all__ = ("Debugger2GrpcAsyncIOTransport",)
8 changes: 7 additions & 1 deletion google/cloud/debugger_v2/types/controller.py
Expand Up @@ -33,6 +33,7 @@

class RegisterDebuggeeRequest(proto.Message):
r"""Request to register a debuggee.

Attributes:
debuggee (google.cloud.debugger_v2.types.Debuggee):
Required. Debuggee information to register. The fields
Expand All @@ -45,6 +46,7 @@ class RegisterDebuggeeRequest(proto.Message):

class RegisterDebuggeeResponse(proto.Message):
r"""Response for registering a debuggee.

Attributes:
debuggee (google.cloud.debugger_v2.types.Debuggee):
Debuggee resource. The field ``id`` is guaranteed to be set
Expand All @@ -60,6 +62,7 @@ class RegisterDebuggeeResponse(proto.Message):

class ListActiveBreakpointsRequest(proto.Message):
r"""Request to list active breakpoints.

Attributes:
debuggee_id (str):
Required. Identifies the debuggee.
Expand Down Expand Up @@ -87,6 +90,7 @@ class ListActiveBreakpointsRequest(proto.Message):

class ListActiveBreakpointsResponse(proto.Message):
r"""Response for listing active breakpoints.

Attributes:
breakpoints (Sequence[google.cloud.debugger_v2.types.Breakpoint]):
List of all active breakpoints. The fields ``id`` and
Expand All @@ -109,6 +113,7 @@ class ListActiveBreakpointsResponse(proto.Message):

class UpdateActiveBreakpointRequest(proto.Message):
r"""Request to update an active breakpoint.

Attributes:
debuggee_id (str):
Required. Identifies the debuggee being
Expand All @@ -126,7 +131,8 @@ class UpdateActiveBreakpointRequest(proto.Message):
class UpdateActiveBreakpointResponse(proto.Message):
r"""Response for updating an active breakpoint.
The message is defined to allow future extensions.
"""

"""


__all__ = tuple(sorted(__protobuf__.manifest))
2 changes: 1 addition & 1 deletion google/cloud/debugger_v2/types/debugger.py
Expand Up @@ -35,7 +35,7 @@


class SetBreakpointRequest(proto.Message):
r"""Request to set a breakpoin.
r"""Request to set a breakpoint

Attributes:
debuggee_id (str):
Expand Down
50 changes: 50 additions & 0 deletions tests/unit/gapic/debugger_v2/test_controller2.py
Expand Up @@ -29,6 +29,7 @@
from google.api_core import gapic_v1
from google.api_core import grpc_helpers
from google.api_core import grpc_helpers_async
from google.api_core import path_template
from google.auth import credentials as ga_credentials
from google.auth.exceptions import MutualTLSChannelError
from google.cloud.debugger_v2.services.controller2 import Controller2AsyncClient
Expand Down Expand Up @@ -1069,6 +1070,9 @@ def test_controller2_base_transport():
with pytest.raises(NotImplementedError):
getattr(transport, method)(request=object())

with pytest.raises(NotImplementedError):
transport.close()


@requires_google_auth_gte_1_25_0
def test_controller2_base_transport_with_credentials_file():
Expand Down Expand Up @@ -1531,3 +1535,49 @@ def test_client_withDEFAULT_CLIENT_INFO():
credentials=ga_credentials.AnonymousCredentials(), client_info=client_info,
)
prep.assert_called_once_with(client_info)


@pytest.mark.asyncio
async def test_transport_close_async():
client = Controller2AsyncClient(
credentials=ga_credentials.AnonymousCredentials(), transport="grpc_asyncio",
)
with mock.patch.object(
type(getattr(client.transport, "grpc_channel")), "close"
) as close:
async with client:
close.assert_not_called()
close.assert_called_once()


def test_transport_close():
transports = {
"grpc": "_grpc_channel",
}

for transport, close_name in transports.items():
client = Controller2Client(
credentials=ga_credentials.AnonymousCredentials(), transport=transport
)
with mock.patch.object(
type(getattr(client.transport, close_name)), "close"
) as close:
with client:
close.assert_not_called()
close.assert_called_once()


def test_client_ctx():
transports = [
"grpc",
]
for transport in transports:
client = Controller2Client(
credentials=ga_credentials.AnonymousCredentials(), transport=transport
)
# Test client calls underlying transport.
with mock.patch.object(type(client.transport), "close") as close:
close.assert_not_called()
with client:
pass
close.assert_called()
50 changes: 50 additions & 0 deletions tests/unit/gapic/debugger_v2/test_debugger2.py
Expand Up @@ -29,6 +29,7 @@
from google.api_core import gapic_v1
from google.api_core import grpc_helpers
from google.api_core import grpc_helpers_async
from google.api_core import path_template
from google.auth import credentials as ga_credentials
from google.auth.exceptions import MutualTLSChannelError
from google.cloud.debugger_v2.services.debugger2 import Debugger2AsyncClient
Expand Down Expand Up @@ -1382,6 +1383,9 @@ def test_debugger2_base_transport():
with pytest.raises(NotImplementedError):
getattr(transport, method)(request=object())

with pytest.raises(NotImplementedError):
transport.close()


@requires_google_auth_gte_1_25_0
def test_debugger2_base_transport_with_credentials_file():
Expand Down Expand Up @@ -1844,3 +1848,49 @@ def test_client_withDEFAULT_CLIENT_INFO():
credentials=ga_credentials.AnonymousCredentials(), client_info=client_info,
)
prep.assert_called_once_with(client_info)


@pytest.mark.asyncio
async def test_transport_close_async():
client = Debugger2AsyncClient(
credentials=ga_credentials.AnonymousCredentials(), transport="grpc_asyncio",
)
with mock.patch.object(
type(getattr(client.transport, "grpc_channel")), "close"
) as close:
async with client:
close.assert_not_called()
close.assert_called_once()


def test_transport_close():
transports = {
"grpc": "_grpc_channel",
}

for transport, close_name in transports.items():
client = Debugger2Client(
credentials=ga_credentials.AnonymousCredentials(), transport=transport
)
with mock.patch.object(
type(getattr(client.transport, close_name)), "close"
) as close:
with client:
close.assert_not_called()
close.assert_called_once()


def test_client_ctx():
transports = [
"grpc",
]
for transport in transports:
client = Debugger2Client(
credentials=ga_credentials.AnonymousCredentials(), transport=transport
)
# Test client calls underlying transport.
with mock.patch.object(type(client.transport), "close") as close:
close.assert_not_called()
with client:
pass
close.assert_called()