feat: Added the VertexAiResourceNoun.to_dict() method (#588)

googleapis · Aug 9, 2021 · b478075 · b478075
1 parent 726dea7
commit b478075
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 6 deletions.
diff --git a/google/cloud/aiplatform/base.py b/google/cloud/aiplatform/base.py
@@ -44,6 +44,7 @@
 from google.cloud.aiplatform import initializer
 from google.cloud.aiplatform import utils
 from google.cloud.aiplatform.compat.types import encryption_spec as gca_encryption_spec
+from google.protobuf import json_format
 
 logging.basicConfig(level=logging.INFO, stream=sys.stdout)
 
@@ -607,6 +608,10 @@ def _assert_gca_resource_is_available(self) -> None:
     def __repr__(self) -> str:
         return f"{object.__repr__(self)} \nresource name: {self.resource_name}"
 
+    def to_dict(self) -> Dict[str, Any]:
+        """Returns the resource proto as a dictionary."""
+        return json_format.MessageToDict(self.gca_resource._pb)
+
 
 def optional_sync(
     construct_object_on_arg: Optional[str] = None,

diff --git a/google/cloud/aiplatform/metadata/resource.py b/google/cloud/aiplatform/metadata/resource.py
@@ -24,7 +24,6 @@
 import proto
 from google.api_core import exceptions
 from google.auth import credentials as auth_credentials
-from google.protobuf import json_format
 
 from google.cloud.aiplatform import base, initializer
 from google.cloud.aiplatform import utils
@@ -98,7 +97,7 @@ def __init__(
 
     @property
     def metadata(self) -> Dict:
-        return json_format.MessageToDict(self._gca_resource._pb)["metadata"]
+        return self.to_dict()["metadata"]
 
     @property
     def schema_title(self) -> str:

diff --git a/tests/system/aiplatform/test_dataset.py b/tests/system/aiplatform/test_dataset.py
@@ -21,7 +21,6 @@
 import importlib
 
 from google import auth as google_auth
-from google.protobuf import json_format
 from google.api_core import exceptions
 from google.api_core import client_options
 
@@ -239,16 +238,15 @@ def test_create_tabular_dataset(self, dataset_gapic_client, shared_state):
             gcs_source=[_TEST_TABULAR_CLASSIFICATION_GCS_SOURCE],
         )
 
-        gapic_dataset = tabular_dataset._gca_resource
         shared_state["dataset_name"] = tabular_dataset.resource_name
 
-        gapic_metadata = json_format.MessageToDict(gapic_dataset._pb.metadata)
+        gapic_metadata = tabular_dataset.to_dict()["metadata"]
         gcs_source_uris = gapic_metadata["inputConfig"]["gcsSource"]["uri"]
 
         assert len(gcs_source_uris) == 1
         assert _TEST_TABULAR_CLASSIFICATION_GCS_SOURCE == gcs_source_uris[0]
         assert (
-            gapic_dataset.metadata_schema_uri
+            tabular_dataset.metadata_schema_uri
             == aiplatform.schema.dataset.metadata.tabular
         )