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

fix: enable self signed jwt for grpc #24

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
4 changes: 4 additions & 0 deletions google/cloud/tpu_v1/services/tpu/client.py
Expand Up @@ -383,6 +383,10 @@ 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")
),
)

def list_nodes(
Expand Down
31 changes: 18 additions & 13 deletions tests/unit/gapic/tpu_v1/test_tpu.py
Expand Up @@ -110,33 +110,28 @@ def test_tpu_client_from_service_account_info(client_class):
assert client.transport._host == "tpu.googleapis.com:443"


@pytest.mark.parametrize("client_class", [TpuClient, TpuAsyncClient,])
def test_tpu_client_service_account_always_use_jwt(client_class):
with mock.patch.object(
service_account.Credentials, "with_always_use_jwt_access", create=True
) as use_jwt:
creds = service_account.Credentials(None, None, None)
client = client_class(credentials=creds)
use_jwt.assert_not_called()


@pytest.mark.parametrize(
"transport_class,transport_name",
[
(transports.TpuGrpcTransport, "grpc"),
(transports.TpuGrpcAsyncIOTransport, "grpc_asyncio"),
],
)
def test_tpu_client_service_account_always_use_jwt_true(
transport_class, transport_name
):
def test_tpu_client_service_account_always_use_jwt(transport_class, transport_name):
with mock.patch.object(
service_account.Credentials, "with_always_use_jwt_access", create=True
) as use_jwt:
creds = service_account.Credentials(None, None, None)
transport = transport_class(credentials=creds, always_use_jwt_access=True)
use_jwt.assert_called_once_with(True)

with mock.patch.object(
service_account.Credentials, "with_always_use_jwt_access", create=True
) as use_jwt:
creds = service_account.Credentials(None, None, None)
transport = transport_class(credentials=creds, always_use_jwt_access=False)
use_jwt.assert_not_called()


@pytest.mark.parametrize("client_class", [TpuClient, TpuAsyncClient,])
def test_tpu_client_from_service_account_file(client_class):
Expand Down Expand Up @@ -203,6 +198,7 @@ def test_tpu_client_client_options(client_class, transport_class, transport_name
client_cert_source_for_mtls=None,
quota_project_id=None,
client_info=transports.base.DEFAULT_CLIENT_INFO,
always_use_jwt_access=True,
)

# Check the case api_endpoint is not provided and GOOGLE_API_USE_MTLS_ENDPOINT is
Expand All @@ -219,6 +215,7 @@ def test_tpu_client_client_options(client_class, transport_class, transport_name
client_cert_source_for_mtls=None,
quota_project_id=None,
client_info=transports.base.DEFAULT_CLIENT_INFO,
always_use_jwt_access=True,
)

# Check the case api_endpoint is not provided and GOOGLE_API_USE_MTLS_ENDPOINT is
Expand All @@ -235,6 +232,7 @@ def test_tpu_client_client_options(client_class, transport_class, transport_name
client_cert_source_for_mtls=None,
quota_project_id=None,
client_info=transports.base.DEFAULT_CLIENT_INFO,
always_use_jwt_access=True,
)

# Check the case api_endpoint is not provided and GOOGLE_API_USE_MTLS_ENDPOINT has
Expand Down Expand Up @@ -263,6 +261,7 @@ def test_tpu_client_client_options(client_class, transport_class, transport_name
client_cert_source_for_mtls=None,
quota_project_id="octopus",
client_info=transports.base.DEFAULT_CLIENT_INFO,
always_use_jwt_access=True,
)


Expand Down Expand Up @@ -313,6 +312,7 @@ def test_tpu_client_mtls_env_auto(
client_cert_source_for_mtls=expected_client_cert_source,
quota_project_id=None,
client_info=transports.base.DEFAULT_CLIENT_INFO,
always_use_jwt_access=True,
)

# Check the case ADC client cert is provided. Whether client cert is used depends on
Expand Down Expand Up @@ -346,6 +346,7 @@ def test_tpu_client_mtls_env_auto(
client_cert_source_for_mtls=expected_client_cert_source,
quota_project_id=None,
client_info=transports.base.DEFAULT_CLIENT_INFO,
always_use_jwt_access=True,
)

# Check the case client_cert_source and ADC client cert are not provided.
Expand All @@ -367,6 +368,7 @@ def test_tpu_client_mtls_env_auto(
client_cert_source_for_mtls=None,
quota_project_id=None,
client_info=transports.base.DEFAULT_CLIENT_INFO,
always_use_jwt_access=True,
)


Expand All @@ -393,6 +395,7 @@ def test_tpu_client_client_options_scopes(
client_cert_source_for_mtls=None,
quota_project_id=None,
client_info=transports.base.DEFAULT_CLIENT_INFO,
always_use_jwt_access=True,
)


Expand All @@ -419,6 +422,7 @@ def test_tpu_client_client_options_credentials_file(
client_cert_source_for_mtls=None,
quota_project_id=None,
client_info=transports.base.DEFAULT_CLIENT_INFO,
always_use_jwt_access=True,
)


Expand All @@ -436,6 +440,7 @@ def test_tpu_client_client_options_from_dict():
client_cert_source_for_mtls=None,
quota_project_id=None,
client_info=transports.base.DEFAULT_CLIENT_INFO,
always_use_jwt_access=True,
)


Expand Down