From 7c97647579c510cd3a37c896914be74759370bfd Mon Sep 17 00:00:00 2001 From: Sasha Sobran Date: Fri, 21 May 2021 09:54:55 -0400 Subject: [PATCH 1/3] fix: use resource name location when passed full resource name --- google/cloud/aiplatform/base.py | 3 --- tests/unit/aiplatform/test_datasets.py | 12 ++++++++++++ 2 files changed, 12 insertions(+), 3 deletions(-) diff --git a/google/cloud/aiplatform/base.py b/google/cloud/aiplatform/base.py index f183c1aedc..07e4c2fe4a 100644 --- a/google/cloud/aiplatform/base.py +++ b/google/cloud/aiplatform/base.py @@ -501,9 +501,6 @@ def _get_and_validate_project_location( RuntimeError if location is different from resource location """ - if not project and not location: - return project, location - fields = utils.extract_fields_from_resource_name( resource_name, self._resource_noun ) diff --git a/tests/unit/aiplatform/test_datasets.py b/tests/unit/aiplatform/test_datasets.py index 6b67d67a20..7307e57ae1 100644 --- a/tests/unit/aiplatform/test_datasets.py +++ b/tests/unit/aiplatform/test_datasets.py @@ -401,6 +401,18 @@ def test_init_dataset_with_alt_project_and_location(self, get_dataset_mock): ) get_dataset_mock.assert_called_once_with(name=_TEST_NAME) + def test_init_dataset_with_alt_location(self, get_dataset_tabular_gcs_mock): + aiplatform.init(project=_TEST_PROJECT, location=_TEST_ALT_LOCATION) + + ds = datasets.TabularDataset(dataset_name=_TEST_NAME) + + assert ds.api_client._clients['v1']._client_options.api_endpoint == f'{_TEST_LOCATION}-{aiplatform.constants.API_BASE_PATH}' + + assert _TEST_ALT_LOCATION != _TEST_LOCATION + + get_dataset_tabular_gcs_mock.assert_called_once_with(name=_TEST_NAME) + + def test_init_dataset_with_project_and_alt_location(self): aiplatform.init(project=_TEST_PROJECT) with pytest.raises(RuntimeError): From ed3a43619078c28a6196656350b394bfc13ab8e6 Mon Sep 17 00:00:00 2001 From: Sasha Sobran Date: Fri, 21 May 2021 10:57:39 -0400 Subject: [PATCH 2/3] chore:lint --- tests/unit/aiplatform/test_datasets.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/unit/aiplatform/test_datasets.py b/tests/unit/aiplatform/test_datasets.py index 7307e57ae1..35f020bedf 100644 --- a/tests/unit/aiplatform/test_datasets.py +++ b/tests/unit/aiplatform/test_datasets.py @@ -406,13 +406,15 @@ def test_init_dataset_with_alt_location(self, get_dataset_tabular_gcs_mock): ds = datasets.TabularDataset(dataset_name=_TEST_NAME) - assert ds.api_client._clients['v1']._client_options.api_endpoint == f'{_TEST_LOCATION}-{aiplatform.constants.API_BASE_PATH}' + assert ( + ds.api_client._clients["v1"]._client_options.api_endpoint + == f"{_TEST_LOCATION}-{aiplatform.constants.API_BASE_PATH}" + ) assert _TEST_ALT_LOCATION != _TEST_LOCATION get_dataset_tabular_gcs_mock.assert_called_once_with(name=_TEST_NAME) - def test_init_dataset_with_project_and_alt_location(self): aiplatform.init(project=_TEST_PROJECT) with pytest.raises(RuntimeError): From 89f54897ca33a84d5f9e6264eec3d2882f6cdced Mon Sep 17 00:00:00 2001 From: Sasha Sobran Date: Fri, 21 May 2021 16:08:34 -0400 Subject: [PATCH 3/3] chore: change tests to use compat --- tests/unit/aiplatform/test_datasets.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/unit/aiplatform/test_datasets.py b/tests/unit/aiplatform/test_datasets.py index 35f020bedf..4c2a75c393 100644 --- a/tests/unit/aiplatform/test_datasets.py +++ b/tests/unit/aiplatform/test_datasets.py @@ -31,6 +31,7 @@ from google.cloud import bigquery from google.cloud import storage +from google.cloud.aiplatform import compat from google.cloud.aiplatform import datasets from google.cloud.aiplatform import initializer from google.cloud.aiplatform import schema @@ -407,7 +408,7 @@ def test_init_dataset_with_alt_location(self, get_dataset_tabular_gcs_mock): ds = datasets.TabularDataset(dataset_name=_TEST_NAME) assert ( - ds.api_client._clients["v1"]._client_options.api_endpoint + ds.api_client._clients[compat.DEFAULT_VERSION]._client_options.api_endpoint == f"{_TEST_LOCATION}-{aiplatform.constants.API_BASE_PATH}" )