diff --git a/google/cloud/dataflow_v1beta3/services/flex_templates_service/async_client.py b/google/cloud/dataflow_v1beta3/services/flex_templates_service/async_client.py index 51e2fe6..4a91023 100644 --- a/google/cloud/dataflow_v1beta3/services/flex_templates_service/async_client.py +++ b/google/cloud/dataflow_v1beta3/services/flex_templates_service/async_client.py @@ -205,6 +205,12 @@ async def launch_flex_template( # 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( diff --git a/google/cloud/dataflow_v1beta3/services/flex_templates_service/client.py b/google/cloud/dataflow_v1beta3/services/flex_templates_service/client.py index c716fea..ca2ee7e 100644 --- a/google/cloud/dataflow_v1beta3/services/flex_templates_service/client.py +++ b/google/cloud/dataflow_v1beta3/services/flex_templates_service/client.py @@ -332,10 +332,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 launch_flex_template( @@ -382,6 +379,19 @@ def launch_flex_template( # 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( diff --git a/google/cloud/dataflow_v1beta3/services/flex_templates_service/transports/base.py b/google/cloud/dataflow_v1beta3/services/flex_templates_service/transports/base.py index 3877a1f..67dccd5 100644 --- a/google/cloud/dataflow_v1beta3/services/flex_templates_service/transports/base.py +++ b/google/cloud/dataflow_v1beta3/services/flex_templates_service/transports/base.py @@ -166,6 +166,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 launch_flex_template( self, diff --git a/google/cloud/dataflow_v1beta3/services/flex_templates_service/transports/grpc.py b/google/cloud/dataflow_v1beta3/services/flex_templates_service/transports/grpc.py index a8a7302..32e2892 100644 --- a/google/cloud/dataflow_v1beta3/services/flex_templates_service/transports/grpc.py +++ b/google/cloud/dataflow_v1beta3/services/flex_templates_service/transports/grpc.py @@ -254,5 +254,8 @@ def launch_flex_template( ) return self._stubs["launch_flex_template"] + def close(self): + self.grpc_channel.close() + __all__ = ("FlexTemplatesServiceGrpcTransport",) diff --git a/google/cloud/dataflow_v1beta3/services/flex_templates_service/transports/grpc_asyncio.py b/google/cloud/dataflow_v1beta3/services/flex_templates_service/transports/grpc_asyncio.py index 5054b54..fafa7d9 100644 --- a/google/cloud/dataflow_v1beta3/services/flex_templates_service/transports/grpc_asyncio.py +++ b/google/cloud/dataflow_v1beta3/services/flex_templates_service/transports/grpc_asyncio.py @@ -258,5 +258,8 @@ def launch_flex_template( ) return self._stubs["launch_flex_template"] + def close(self): + return self.grpc_channel.close() + __all__ = ("FlexTemplatesServiceGrpcAsyncIOTransport",) diff --git a/google/cloud/dataflow_v1beta3/services/jobs_v1_beta3/async_client.py b/google/cloud/dataflow_v1beta3/services/jobs_v1_beta3/async_client.py index 1386ba0..90a555f 100644 --- a/google/cloud/dataflow_v1beta3/services/jobs_v1_beta3/async_client.py +++ b/google/cloud/dataflow_v1beta3/services/jobs_v1_beta3/async_client.py @@ -517,6 +517,12 @@ async def snapshot_job( # 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( diff --git a/google/cloud/dataflow_v1beta3/services/jobs_v1_beta3/client.py b/google/cloud/dataflow_v1beta3/services/jobs_v1_beta3/client.py index ece7b4d..120a1ba 100644 --- a/google/cloud/dataflow_v1beta3/services/jobs_v1_beta3/client.py +++ b/google/cloud/dataflow_v1beta3/services/jobs_v1_beta3/client.py @@ -333,10 +333,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 create_job( @@ -701,6 +698,19 @@ def snapshot_job( # 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( diff --git a/google/cloud/dataflow_v1beta3/services/jobs_v1_beta3/transports/base.py b/google/cloud/dataflow_v1beta3/services/jobs_v1_beta3/transports/base.py index 472ea83..aa6f2e3 100644 --- a/google/cloud/dataflow_v1beta3/services/jobs_v1_beta3/transports/base.py +++ b/google/cloud/dataflow_v1beta3/services/jobs_v1_beta3/transports/base.py @@ -185,6 +185,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 create_job( self, diff --git a/google/cloud/dataflow_v1beta3/services/jobs_v1_beta3/transports/grpc.py b/google/cloud/dataflow_v1beta3/services/jobs_v1_beta3/transports/grpc.py index 2c37c26..e522a7d 100644 --- a/google/cloud/dataflow_v1beta3/services/jobs_v1_beta3/transports/grpc.py +++ b/google/cloud/dataflow_v1beta3/services/jobs_v1_beta3/transports/grpc.py @@ -428,5 +428,8 @@ def snapshot_job(self) -> Callable[[jobs.SnapshotJobRequest], snapshots.Snapshot ) return self._stubs["snapshot_job"] + def close(self): + self.grpc_channel.close() + __all__ = ("JobsV1Beta3GrpcTransport",) diff --git a/google/cloud/dataflow_v1beta3/services/jobs_v1_beta3/transports/grpc_asyncio.py b/google/cloud/dataflow_v1beta3/services/jobs_v1_beta3/transports/grpc_asyncio.py index 524eb04..4326920 100644 --- a/google/cloud/dataflow_v1beta3/services/jobs_v1_beta3/transports/grpc_asyncio.py +++ b/google/cloud/dataflow_v1beta3/services/jobs_v1_beta3/transports/grpc_asyncio.py @@ -437,5 +437,8 @@ def snapshot_job( ) return self._stubs["snapshot_job"] + def close(self): + return self.grpc_channel.close() + __all__ = ("JobsV1Beta3GrpcAsyncIOTransport",) diff --git a/google/cloud/dataflow_v1beta3/services/messages_v1_beta3/async_client.py b/google/cloud/dataflow_v1beta3/services/messages_v1_beta3/async_client.py index 3b78e74..bdb2182 100644 --- a/google/cloud/dataflow_v1beta3/services/messages_v1_beta3/async_client.py +++ b/google/cloud/dataflow_v1beta3/services/messages_v1_beta3/async_client.py @@ -224,6 +224,12 @@ async def list_job_messages( # 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( diff --git a/google/cloud/dataflow_v1beta3/services/messages_v1_beta3/client.py b/google/cloud/dataflow_v1beta3/services/messages_v1_beta3/client.py index 4c4bec4..0e91ff8 100644 --- a/google/cloud/dataflow_v1beta3/services/messages_v1_beta3/client.py +++ b/google/cloud/dataflow_v1beta3/services/messages_v1_beta3/client.py @@ -330,10 +330,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 list_job_messages( @@ -400,6 +397,19 @@ def list_job_messages( # 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( diff --git a/google/cloud/dataflow_v1beta3/services/messages_v1_beta3/transports/base.py b/google/cloud/dataflow_v1beta3/services/messages_v1_beta3/transports/base.py index 5812189..5f3e28e 100644 --- a/google/cloud/dataflow_v1beta3/services/messages_v1_beta3/transports/base.py +++ b/google/cloud/dataflow_v1beta3/services/messages_v1_beta3/transports/base.py @@ -164,6 +164,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 list_job_messages( self, diff --git a/google/cloud/dataflow_v1beta3/services/messages_v1_beta3/transports/grpc.py b/google/cloud/dataflow_v1beta3/services/messages_v1_beta3/transports/grpc.py index aa4b6d5..0afbe02 100644 --- a/google/cloud/dataflow_v1beta3/services/messages_v1_beta3/transports/grpc.py +++ b/google/cloud/dataflow_v1beta3/services/messages_v1_beta3/transports/grpc.py @@ -260,5 +260,8 @@ def list_job_messages( ) return self._stubs["list_job_messages"] + def close(self): + self.grpc_channel.close() + __all__ = ("MessagesV1Beta3GrpcTransport",) diff --git a/google/cloud/dataflow_v1beta3/services/messages_v1_beta3/transports/grpc_asyncio.py b/google/cloud/dataflow_v1beta3/services/messages_v1_beta3/transports/grpc_asyncio.py index 5b6b71e..d1213af 100644 --- a/google/cloud/dataflow_v1beta3/services/messages_v1_beta3/transports/grpc_asyncio.py +++ b/google/cloud/dataflow_v1beta3/services/messages_v1_beta3/transports/grpc_asyncio.py @@ -265,5 +265,8 @@ def list_job_messages( ) return self._stubs["list_job_messages"] + def close(self): + return self.grpc_channel.close() + __all__ = ("MessagesV1Beta3GrpcAsyncIOTransport",) diff --git a/google/cloud/dataflow_v1beta3/services/metrics_v1_beta3/async_client.py b/google/cloud/dataflow_v1beta3/services/metrics_v1_beta3/async_client.py index 91ef9a4..7ac8193 100644 --- a/google/cloud/dataflow_v1beta3/services/metrics_v1_beta3/async_client.py +++ b/google/cloud/dataflow_v1beta3/services/metrics_v1_beta3/async_client.py @@ -333,6 +333,12 @@ async def get_stage_execution_details( # 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( diff --git a/google/cloud/dataflow_v1beta3/services/metrics_v1_beta3/client.py b/google/cloud/dataflow_v1beta3/services/metrics_v1_beta3/client.py index b0951ee..eb1e4ff 100644 --- a/google/cloud/dataflow_v1beta3/services/metrics_v1_beta3/client.py +++ b/google/cloud/dataflow_v1beta3/services/metrics_v1_beta3/client.py @@ -331,10 +331,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 get_job_metrics( @@ -515,6 +512,19 @@ def get_stage_execution_details( # 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( diff --git a/google/cloud/dataflow_v1beta3/services/metrics_v1_beta3/transports/base.py b/google/cloud/dataflow_v1beta3/services/metrics_v1_beta3/transports/base.py index f522c1c..72fcf85 100644 --- a/google/cloud/dataflow_v1beta3/services/metrics_v1_beta3/transports/base.py +++ b/google/cloud/dataflow_v1beta3/services/metrics_v1_beta3/transports/base.py @@ -174,6 +174,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 get_job_metrics( self, diff --git a/google/cloud/dataflow_v1beta3/services/metrics_v1_beta3/transports/grpc.py b/google/cloud/dataflow_v1beta3/services/metrics_v1_beta3/transports/grpc.py index 5043d93..7c01ca7 100644 --- a/google/cloud/dataflow_v1beta3/services/metrics_v1_beta3/transports/grpc.py +++ b/google/cloud/dataflow_v1beta3/services/metrics_v1_beta3/transports/grpc.py @@ -321,5 +321,8 @@ def get_stage_execution_details( ) return self._stubs["get_stage_execution_details"] + def close(self): + self.grpc_channel.close() + __all__ = ("MetricsV1Beta3GrpcTransport",) diff --git a/google/cloud/dataflow_v1beta3/services/metrics_v1_beta3/transports/grpc_asyncio.py b/google/cloud/dataflow_v1beta3/services/metrics_v1_beta3/transports/grpc_asyncio.py index 9a032fa..58b094d 100644 --- a/google/cloud/dataflow_v1beta3/services/metrics_v1_beta3/transports/grpc_asyncio.py +++ b/google/cloud/dataflow_v1beta3/services/metrics_v1_beta3/transports/grpc_asyncio.py @@ -327,5 +327,8 @@ def get_stage_execution_details( ) return self._stubs["get_stage_execution_details"] + def close(self): + return self.grpc_channel.close() + __all__ = ("MetricsV1Beta3GrpcAsyncIOTransport",) diff --git a/google/cloud/dataflow_v1beta3/services/snapshots_v1_beta3/async_client.py b/google/cloud/dataflow_v1beta3/services/snapshots_v1_beta3/async_client.py index bb59b5c..4fd19d2 100644 --- a/google/cloud/dataflow_v1beta3/services/snapshots_v1_beta3/async_client.py +++ b/google/cloud/dataflow_v1beta3/services/snapshots_v1_beta3/async_client.py @@ -283,6 +283,12 @@ async def list_snapshots( # 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( diff --git a/google/cloud/dataflow_v1beta3/services/snapshots_v1_beta3/client.py b/google/cloud/dataflow_v1beta3/services/snapshots_v1_beta3/client.py index 2c6e592..237af6a 100644 --- a/google/cloud/dataflow_v1beta3/services/snapshots_v1_beta3/client.py +++ b/google/cloud/dataflow_v1beta3/services/snapshots_v1_beta3/client.py @@ -331,10 +331,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 get_snapshot( @@ -461,6 +458,19 @@ def list_snapshots( # 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( diff --git a/google/cloud/dataflow_v1beta3/services/snapshots_v1_beta3/transports/base.py b/google/cloud/dataflow_v1beta3/services/snapshots_v1_beta3/transports/base.py index 2d869c1..55dbee2 100644 --- a/google/cloud/dataflow_v1beta3/services/snapshots_v1_beta3/transports/base.py +++ b/google/cloud/dataflow_v1beta3/services/snapshots_v1_beta3/transports/base.py @@ -170,6 +170,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 get_snapshot( self, diff --git a/google/cloud/dataflow_v1beta3/services/snapshots_v1_beta3/transports/grpc.py b/google/cloud/dataflow_v1beta3/services/snapshots_v1_beta3/transports/grpc.py index e235b96..2c32a08 100644 --- a/google/cloud/dataflow_v1beta3/services/snapshots_v1_beta3/transports/grpc.py +++ b/google/cloud/dataflow_v1beta3/services/snapshots_v1_beta3/transports/grpc.py @@ -304,5 +304,8 @@ def list_snapshots( ) return self._stubs["list_snapshots"] + def close(self): + self.grpc_channel.close() + __all__ = ("SnapshotsV1Beta3GrpcTransport",) diff --git a/google/cloud/dataflow_v1beta3/services/snapshots_v1_beta3/transports/grpc_asyncio.py b/google/cloud/dataflow_v1beta3/services/snapshots_v1_beta3/transports/grpc_asyncio.py index e5304ee..2731cdf 100644 --- a/google/cloud/dataflow_v1beta3/services/snapshots_v1_beta3/transports/grpc_asyncio.py +++ b/google/cloud/dataflow_v1beta3/services/snapshots_v1_beta3/transports/grpc_asyncio.py @@ -311,5 +311,8 @@ def list_snapshots( ) return self._stubs["list_snapshots"] + def close(self): + return self.grpc_channel.close() + __all__ = ("SnapshotsV1Beta3GrpcAsyncIOTransport",) diff --git a/google/cloud/dataflow_v1beta3/services/templates_service/async_client.py b/google/cloud/dataflow_v1beta3/services/templates_service/async_client.py index 071ec3a..ed37177 100644 --- a/google/cloud/dataflow_v1beta3/services/templates_service/async_client.py +++ b/google/cloud/dataflow_v1beta3/services/templates_service/async_client.py @@ -292,6 +292,12 @@ async def get_template( # 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( diff --git a/google/cloud/dataflow_v1beta3/services/templates_service/client.py b/google/cloud/dataflow_v1beta3/services/templates_service/client.py index d258165..89c8fcc 100644 --- a/google/cloud/dataflow_v1beta3/services/templates_service/client.py +++ b/google/cloud/dataflow_v1beta3/services/templates_service/client.py @@ -333,10 +333,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 create_job_from_template( @@ -470,6 +467,19 @@ def get_template( # 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( diff --git a/google/cloud/dataflow_v1beta3/services/templates_service/transports/base.py b/google/cloud/dataflow_v1beta3/services/templates_service/transports/base.py index eb3a928..c12f0b9 100644 --- a/google/cloud/dataflow_v1beta3/services/templates_service/transports/base.py +++ b/google/cloud/dataflow_v1beta3/services/templates_service/transports/base.py @@ -173,6 +173,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 create_job_from_template( self, diff --git a/google/cloud/dataflow_v1beta3/services/templates_service/transports/grpc.py b/google/cloud/dataflow_v1beta3/services/templates_service/transports/grpc.py index 8b0308d..729ba13 100644 --- a/google/cloud/dataflow_v1beta3/services/templates_service/transports/grpc.py +++ b/google/cloud/dataflow_v1beta3/services/templates_service/transports/grpc.py @@ -305,5 +305,8 @@ def get_template( ) return self._stubs["get_template"] + def close(self): + self.grpc_channel.close() + __all__ = ("TemplatesServiceGrpcTransport",) diff --git a/google/cloud/dataflow_v1beta3/services/templates_service/transports/grpc_asyncio.py b/google/cloud/dataflow_v1beta3/services/templates_service/transports/grpc_asyncio.py index 405a521..c0da99c 100644 --- a/google/cloud/dataflow_v1beta3/services/templates_service/transports/grpc_asyncio.py +++ b/google/cloud/dataflow_v1beta3/services/templates_service/transports/grpc_asyncio.py @@ -312,5 +312,8 @@ def get_template( ) return self._stubs["get_template"] + def close(self): + return self.grpc_channel.close() + __all__ = ("TemplatesServiceGrpcAsyncIOTransport",) diff --git a/google/cloud/dataflow_v1beta3/types/environment.py b/google/cloud/dataflow_v1beta3/types/environment.py index 98fe1a7..2731689 100644 --- a/google/cloud/dataflow_v1beta3/types/environment.py +++ b/google/cloud/dataflow_v1beta3/types/environment.py @@ -116,6 +116,7 @@ class ShuffleMode(proto.Enum): class Environment(proto.Message): r"""Describes the environment in which a Dataflow Job runs. + Attributes: temp_storage_prefix (str): The prefix of the resources the system should use for @@ -264,6 +265,7 @@ class Package(proto.Message): class Disk(proto.Message): r"""Describes the data disk used by a workflow job. + Attributes: size_gb (int): Size of disk in GB. If zero or unspecified, @@ -305,6 +307,7 @@ class Disk(proto.Message): class WorkerSettings(proto.Message): r"""Provides data to pass through to the worker harness. + Attributes: base_url (str): The base URL for accessing Google Cloud APIs. @@ -351,6 +354,7 @@ class WorkerSettings(proto.Message): class TaskRunnerSettings(proto.Message): r"""Taskrunner configuration settings. + Attributes: task_user (str): The UNIX user ID on the worker VM to use for @@ -450,6 +454,7 @@ class TaskRunnerSettings(proto.Message): class AutoscalingSettings(proto.Message): r"""Settings for WorkerPool autoscaling. + Attributes: algorithm (google.cloud.dataflow_v1beta3.types.AutoscalingAlgorithm): The algorithm to use for autoscaling. diff --git a/google/cloud/dataflow_v1beta3/types/jobs.py b/google/cloud/dataflow_v1beta3/types/jobs.py index 12d9683..5180087 100644 --- a/google/cloud/dataflow_v1beta3/types/jobs.py +++ b/google/cloud/dataflow_v1beta3/types/jobs.py @@ -292,6 +292,7 @@ class Job(proto.Message): class DatastoreIODetails(proto.Message): r"""Metadata for a Datastore connector used by the job. + Attributes: namespace (str): Namespace used in the connection. @@ -305,6 +306,7 @@ class DatastoreIODetails(proto.Message): class PubSubIODetails(proto.Message): r"""Metadata for a Pub/Sub connector used by the job. + Attributes: topic (str): Topic accessed in the connection. @@ -318,6 +320,7 @@ class PubSubIODetails(proto.Message): class FileIODetails(proto.Message): r"""Metadata for a File connector used by the job. + Attributes: file_pattern (str): File Pattern used to access files by the @@ -329,6 +332,7 @@ class FileIODetails(proto.Message): class BigTableIODetails(proto.Message): r"""Metadata for a Cloud BigTable connector used by the job. + Attributes: project_id (str): ProjectId accessed in the connection. @@ -345,6 +349,7 @@ class BigTableIODetails(proto.Message): class BigQueryIODetails(proto.Message): r"""Metadata for a BigQuery connector used by the job. + Attributes: table (str): Table accessed in the connection. @@ -364,6 +369,7 @@ class BigQueryIODetails(proto.Message): class SpannerIODetails(proto.Message): r"""Metadata for a Spanner connector used by the job. + Attributes: project_id (str): ProjectId accessed in the connection. @@ -380,6 +386,7 @@ class SpannerIODetails(proto.Message): class SdkVersion(proto.Message): r"""The version of the SDK used to run the job. + Attributes: version (str): The version of the SDK used to run the job. @@ -560,6 +567,7 @@ class ExecutionStageSummary(proto.Message): class StageSource(proto.Message): r"""Description of an input or output of an execution stage. + Attributes: user_name (str): Human-readable name for this source; may be @@ -781,6 +789,7 @@ class JobExecutionStageInfo(proto.Message): class CreateJobRequest(proto.Message): r"""Request to create a Cloud Dataflow job. + Attributes: project_id (str): The ID of the Cloud Platform project that the @@ -808,6 +817,7 @@ class CreateJobRequest(proto.Message): class GetJobRequest(proto.Message): r"""Request to get the state of a Cloud Dataflow job. + Attributes: project_id (str): The ID of the Cloud Platform project that the @@ -831,6 +841,7 @@ class GetJobRequest(proto.Message): class UpdateJobRequest(proto.Message): r"""Request to update a Cloud Dataflow job. + Attributes: project_id (str): The ID of the Cloud Platform project that the @@ -855,6 +866,7 @@ class UpdateJobRequest(proto.Message): class ListJobsRequest(proto.Message): r"""Request to list Cloud Dataflow jobs. + Attributes: filter (google.cloud.dataflow_v1beta3.types.ListJobsRequest.Filter): The kind of filter to use. @@ -941,6 +953,7 @@ def raw_page(self): class SnapshotJobRequest(proto.Message): r"""Request to create a snapshot of a job. + Attributes: project_id (str): The project which owns the job to be @@ -969,6 +982,7 @@ class SnapshotJobRequest(proto.Message): class CheckActiveJobsRequest(proto.Message): r"""Request to check is active jobs exists for a project + Attributes: project_id (str): The project which owns the jobs. @@ -979,6 +993,7 @@ class CheckActiveJobsRequest(proto.Message): class CheckActiveJobsResponse(proto.Message): r"""Response for CheckActiveJobsRequest. + Attributes: active_jobs_exist (bool): If True, active jobs exists for project. diff --git a/google/cloud/dataflow_v1beta3/types/messages.py b/google/cloud/dataflow_v1beta3/types/messages.py index a3dca45..355ab79 100644 --- a/google/cloud/dataflow_v1beta3/types/messages.py +++ b/google/cloud/dataflow_v1beta3/types/messages.py @@ -44,6 +44,7 @@ class JobMessageImportance(proto.Enum): class JobMessage(proto.Message): r"""A particular message pertaining to a Dataflow job. + Attributes: id (str): Deprecated. @@ -80,6 +81,7 @@ class StructuredMessage(proto.Message): class Parameter(proto.Message): r"""Structured data associated with this message. + Attributes: key (str): Key or name for this parameter. @@ -186,6 +188,7 @@ class ListJobMessagesRequest(proto.Message): class ListJobMessagesResponse(proto.Message): r"""Response to a request to list job messages. + Attributes: job_messages (Sequence[google.cloud.dataflow_v1beta3.types.JobMessage]): Messages in ascending timestamp order. diff --git a/google/cloud/dataflow_v1beta3/types/metrics.py b/google/cloud/dataflow_v1beta3/types/metrics.py index 85b2e22..887470f 100644 --- a/google/cloud/dataflow_v1beta3/types/metrics.py +++ b/google/cloud/dataflow_v1beta3/types/metrics.py @@ -78,6 +78,7 @@ class MetricStructuredName(proto.Message): class MetricUpdate(proto.Message): r"""Describes the state of a metric. + Attributes: name (google.cloud.dataflow_v1beta3.types.MetricStructuredName): Name of the metric. @@ -152,6 +153,7 @@ class MetricUpdate(proto.Message): class GetJobMetricsRequest(proto.Message): r"""Request to get job metrics. + Attributes: project_id (str): A project id. @@ -196,6 +198,7 @@ class JobMetrics(proto.Message): class GetJobExecutionDetailsRequest(proto.Message): r"""Request to get job execution details. + Attributes: project_id (str): A project id. @@ -238,6 +241,7 @@ class ProgressTimeseries(proto.Message): class Point(proto.Message): r"""A point in the timeseries. + Attributes: time (google.protobuf.timestamp_pb2.Timestamp): The timestamp of the point. @@ -254,6 +258,7 @@ class Point(proto.Message): class StageSummary(proto.Message): r"""Information about a particular execution stage of a job. + Attributes: stage_id (str): ID of this stage @@ -283,6 +288,7 @@ class StageSummary(proto.Message): class JobExecutionDetails(proto.Message): r"""Information about the execution of a job. + Attributes: stages (Sequence[google.cloud.dataflow_v1beta3.types.StageSummary]): The stages of the job execution. @@ -345,6 +351,7 @@ class GetStageExecutionDetailsRequest(proto.Message): class WorkItemDetails(proto.Message): r"""Information about an individual work item execution. + Attributes: task_id (str): Name of this work item. @@ -376,6 +383,7 @@ class WorkItemDetails(proto.Message): class WorkerDetails(proto.Message): r"""Information about a worker + Attributes: worker_name (str): Name of this worker @@ -392,6 +400,7 @@ class WorkerDetails(proto.Message): class StageExecutionDetails(proto.Message): r"""Information about the workers and work items within a stage. + Attributes: workers (Sequence[google.cloud.dataflow_v1beta3.types.WorkerDetails]): Workers that have done work on the stage. diff --git a/google/cloud/dataflow_v1beta3/types/snapshots.py b/google/cloud/dataflow_v1beta3/types/snapshots.py index ebe7b79..9062c58 100644 --- a/google/cloud/dataflow_v1beta3/types/snapshots.py +++ b/google/cloud/dataflow_v1beta3/types/snapshots.py @@ -46,6 +46,7 @@ class SnapshotState(proto.Enum): class PubsubSnapshotMetadata(proto.Message): r"""Represents a Pubsub snapshot. + Attributes: topic_name (str): The name of the Pubsub topic. @@ -62,6 +63,7 @@ class PubsubSnapshotMetadata(proto.Message): class Snapshot(proto.Message): r"""Represents a snapshot of a job. + Attributes: id (str): The unique ID of this snapshot. @@ -107,6 +109,7 @@ class Snapshot(proto.Message): class GetSnapshotRequest(proto.Message): r"""Request to get information about a snapshot + Attributes: project_id (str): The ID of the Cloud Platform project that the @@ -124,6 +127,7 @@ class GetSnapshotRequest(proto.Message): class DeleteSnapshotRequest(proto.Message): r"""Request to delete a snapshot. + Attributes: project_id (str): The ID of the Cloud Platform project that the @@ -140,11 +144,13 @@ class DeleteSnapshotRequest(proto.Message): class DeleteSnapshotResponse(proto.Message): - r"""Response from deleting a snapshot. """ + r"""Response from deleting a snapshot. + """ class ListSnapshotsRequest(proto.Message): r"""Request to list snapshots. + Attributes: project_id (str): The project ID to list snapshots for. @@ -162,6 +168,7 @@ class ListSnapshotsRequest(proto.Message): class ListSnapshotsResponse(proto.Message): r"""List of snapshots. + Attributes: snapshots (Sequence[google.cloud.dataflow_v1beta3.types.Snapshot]): Returned snapshots. diff --git a/google/cloud/dataflow_v1beta3/types/streaming.py b/google/cloud/dataflow_v1beta3/types/streaming.py index 255a62e..927a574 100644 --- a/google/cloud/dataflow_v1beta3/types/streaming.py +++ b/google/cloud/dataflow_v1beta3/types/streaming.py @@ -128,6 +128,7 @@ class StreamingStageLocation(proto.Message): class StreamingSideInputLocation(proto.Message): r"""Identifies the location of a streaming side input. + Attributes: tag (str): Identifies the particular side input within @@ -143,6 +144,7 @@ class StreamingSideInputLocation(proto.Message): class CustomSourceLocation(proto.Message): r"""Identifies the location of a custom souce. + Attributes: stateful (bool): Whether this source is stateful. @@ -183,6 +185,7 @@ class StreamLocation(proto.Message): class StateFamilyConfig(proto.Message): r"""State family configuration. + Attributes: state_family (str): The state family value. @@ -197,6 +200,7 @@ class StateFamilyConfig(proto.Message): class ComputationTopology(proto.Message): r"""All configuration data for a particular Computation. + Attributes: system_stage_name (str): The system stage name. @@ -260,6 +264,7 @@ class KeyRangeLocation(proto.Message): class MountedDataDisk(proto.Message): r"""Describes mounted data disk. + Attributes: data_disk (str): The name of the data disk. @@ -274,6 +279,7 @@ class MountedDataDisk(proto.Message): class DataDiskAssignment(proto.Message): r"""Data disk assignment for a given VM instance. + Attributes: vm_instance (str): VM instance name the data disks mounted to, @@ -338,6 +344,7 @@ class StreamingComputationRanges(proto.Message): class StreamingApplianceSnapshotConfig(proto.Message): r"""Streaming appliance snapshot configuration. + Attributes: snapshot_id (str): If set, indicates the snapshot id for the diff --git a/google/cloud/dataflow_v1beta3/types/templates.py b/google/cloud/dataflow_v1beta3/types/templates.py index c8b4e62..400c52c 100644 --- a/google/cloud/dataflow_v1beta3/types/templates.py +++ b/google/cloud/dataflow_v1beta3/types/templates.py @@ -64,6 +64,7 @@ class ParameterType(proto.Enum): class LaunchFlexTemplateResponse(proto.Message): r"""Response to the request to launch a job from Flex Template. + Attributes: job (google.cloud.dataflow_v1beta3.types.Job): The job that was launched, if the request was @@ -76,6 +77,7 @@ class LaunchFlexTemplateResponse(proto.Message): class ContainerSpec(proto.Message): r"""Container Spec. + Attributes: image (str): Name of the docker container image. E.g., @@ -99,6 +101,7 @@ class ContainerSpec(proto.Message): class LaunchFlexTemplateParameter(proto.Message): r"""Launch FlexTemplate Parameter. + Attributes: job_name (str): Required. The job name to use for the created @@ -253,6 +256,7 @@ class FlexTemplateRuntimeEnvironment(proto.Message): class LaunchFlexTemplateRequest(proto.Message): r"""A request to launch a Cloud Dataflow job from a FlexTemplate. + Attributes: project_id (str): Required. The ID of the Cloud Platform @@ -279,6 +283,7 @@ class LaunchFlexTemplateRequest(proto.Message): class RuntimeEnvironment(proto.Message): r"""The environment values to set at runtime. + Attributes: num_workers (int): The initial number of Google Compute Engine @@ -374,6 +379,7 @@ class RuntimeEnvironment(proto.Message): class ParameterMetadata(proto.Message): r"""Metadata for a specific parameter. + Attributes: name (str): Required. The name of the parameter. @@ -408,6 +414,7 @@ class ParameterMetadata(proto.Message): class TemplateMetadata(proto.Message): r"""Metadata describing a template. + Attributes: name (str): Required. The name of the template. @@ -426,6 +433,7 @@ class TemplateMetadata(proto.Message): class SDKInfo(proto.Message): r"""SDK Information. + Attributes: language (google.cloud.dataflow_v1beta3.types.SDKInfo.Language): Required. The SDK Language. @@ -445,6 +453,7 @@ class Language(proto.Enum): class RuntimeMetadata(proto.Message): r"""RuntimeMetadata describing a runtime environment. + Attributes: sdk_info (google.cloud.dataflow_v1beta3.types.SDKInfo): SDK Info for the template. @@ -460,6 +469,7 @@ class RuntimeMetadata(proto.Message): class CreateJobFromTemplateRequest(proto.Message): r"""A request to create a Cloud Dataflow job from a template. + Attributes: project_id (str): Required. The ID of the Cloud Platform @@ -491,6 +501,7 @@ class CreateJobFromTemplateRequest(proto.Message): class GetTemplateRequest(proto.Message): r"""A request to retrieve a Cloud Dataflow job template. + Attributes: project_id (str): Required. The ID of the Cloud Platform @@ -520,6 +531,7 @@ class TemplateView(proto.Enum): class GetTemplateResponse(proto.Message): r"""The response to a GetTemplate request. + Attributes: status (google.rpc.status_pb2.Status): The status of the get template request. Any problems with @@ -548,6 +560,7 @@ class TemplateType(proto.Enum): class LaunchTemplateParameters(proto.Message): r"""Parameters to provide to the template being launched. + Attributes: job_name (str): Required. The job name to use for the created @@ -576,6 +589,7 @@ class LaunchTemplateParameters(proto.Message): class LaunchTemplateRequest(proto.Message): r"""A request to launch a template. + Attributes: project_id (str): Required. The ID of the Cloud Platform @@ -617,6 +631,7 @@ class LaunchTemplateRequest(proto.Message): class LaunchTemplateResponse(proto.Message): r"""Response to the request to launch a template. + Attributes: job (google.cloud.dataflow_v1beta3.types.Job): The job that was launched, if the request was @@ -639,6 +654,7 @@ class InvalidTemplateParameters(proto.Message): class ParameterViolation(proto.Message): r"""A specific template-parameter violation. + Attributes: parameter (str): The parameter that failed to validate. diff --git a/tests/unit/gapic/dataflow_v1beta3/test_flex_templates_service.py b/tests/unit/gapic/dataflow_v1beta3/test_flex_templates_service.py index 199d3eb..41606f6 100644 --- a/tests/unit/gapic/dataflow_v1beta3/test_flex_templates_service.py +++ b/tests/unit/gapic/dataflow_v1beta3/test_flex_templates_service.py @@ -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.dataflow_v1beta3.services.flex_templates_service import ( @@ -710,6 +711,9 @@ def test_flex_templates_service_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_flex_templates_service_base_transport_with_credentials_file(): @@ -1205,3 +1209,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 = FlexTemplatesServiceAsyncClient( + 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 = FlexTemplatesServiceClient( + 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 = FlexTemplatesServiceClient( + 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() diff --git a/tests/unit/gapic/dataflow_v1beta3/test_jobs_v1_beta3.py b/tests/unit/gapic/dataflow_v1beta3/test_jobs_v1_beta3.py index a77ee69..47a7d0e 100644 --- a/tests/unit/gapic/dataflow_v1beta3/test_jobs_v1_beta3.py +++ b/tests/unit/gapic/dataflow_v1beta3/test_jobs_v1_beta3.py @@ -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.dataflow_v1beta3.services.jobs_v1_beta3 import JobsV1Beta3AsyncClient @@ -1568,6 +1569,9 @@ def test_jobs_v1_beta3_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_jobs_v1_beta3_base_transport_with_credentials_file(): @@ -2044,3 +2048,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 = JobsV1Beta3AsyncClient( + 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 = JobsV1Beta3Client( + 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 = JobsV1Beta3Client( + 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() diff --git a/tests/unit/gapic/dataflow_v1beta3/test_messages_v1_beta3.py b/tests/unit/gapic/dataflow_v1beta3/test_messages_v1_beta3.py index 24547fe..f0e8be8 100644 --- a/tests/unit/gapic/dataflow_v1beta3/test_messages_v1_beta3.py +++ b/tests/unit/gapic/dataflow_v1beta3/test_messages_v1_beta3.py @@ -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.dataflow_v1beta3.services.messages_v1_beta3 import ( @@ -846,6 +847,9 @@ def test_messages_v1_beta3_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_messages_v1_beta3_base_transport_with_credentials_file(): @@ -1339,3 +1343,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 = MessagesV1Beta3AsyncClient( + 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 = MessagesV1Beta3Client( + 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 = MessagesV1Beta3Client( + 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() diff --git a/tests/unit/gapic/dataflow_v1beta3/test_metrics_v1_beta3.py b/tests/unit/gapic/dataflow_v1beta3/test_metrics_v1_beta3.py index d2bf1ef..1e911a5 100644 --- a/tests/unit/gapic/dataflow_v1beta3/test_metrics_v1_beta3.py +++ b/tests/unit/gapic/dataflow_v1beta3/test_metrics_v1_beta3.py @@ -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.dataflow_v1beta3.services.metrics_v1_beta3 import ( @@ -1144,6 +1145,9 @@ def test_metrics_v1_beta3_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_metrics_v1_beta3_base_transport_with_credentials_file(): @@ -1637,3 +1641,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 = MetricsV1Beta3AsyncClient( + 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 = MetricsV1Beta3Client( + 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 = MetricsV1Beta3Client( + 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() diff --git a/tests/unit/gapic/dataflow_v1beta3/test_snapshots_v1_beta3.py b/tests/unit/gapic/dataflow_v1beta3/test_snapshots_v1_beta3.py index ed2239f..b25639f 100644 --- a/tests/unit/gapic/dataflow_v1beta3/test_snapshots_v1_beta3.py +++ b/tests/unit/gapic/dataflow_v1beta3/test_snapshots_v1_beta3.py @@ -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.dataflow_v1beta3.services.snapshots_v1_beta3 import ( @@ -882,6 +883,9 @@ def test_snapshots_v1_beta3_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_snapshots_v1_beta3_base_transport_with_credentials_file(): @@ -1375,3 +1379,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 = SnapshotsV1Beta3AsyncClient( + 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 = SnapshotsV1Beta3Client( + 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 = SnapshotsV1Beta3Client( + 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() diff --git a/tests/unit/gapic/dataflow_v1beta3/test_templates_service.py b/tests/unit/gapic/dataflow_v1beta3/test_templates_service.py index 348d1a7..26dcf1f 100644 --- a/tests/unit/gapic/dataflow_v1beta3/test_templates_service.py +++ b/tests/unit/gapic/dataflow_v1beta3/test_templates_service.py @@ -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.dataflow_v1beta3.services.templates_service import ( @@ -924,6 +925,9 @@ def test_templates_service_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_templates_service_base_transport_with_credentials_file(): @@ -1417,3 +1421,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 = TemplatesServiceAsyncClient( + 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 = TemplatesServiceClient( + 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 = TemplatesServiceClient( + 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()