chore: check default MetadataStore when creating Artifacts

jaycee-li · copybara-github · commit 65300c43bb9f · 2022-12-28T10:58:49.000-08:00
PiperOrigin-RevId: 498214023
diff --git a/google/cloud/aiplatform/metadata/artifact.py b/google/cloud/aiplatform/metadata/artifact.py
@@ -344,6 +344,11 @@ def create(
                 "aiplatform.metadata.artifact.Artifact.create"
             )
 
+        if metadata_store_id == "default":
+            metadata_store._MetadataStore.ensure_default_metadata_store_exists(
+                project=project, location=location, credentials=credentials
+            )
+
         return cls._create(
             resource_id=resource_id,
             schema_title=schema_title,
diff --git a/tests/unit/aiplatform/test_metadata_resources.py b/tests/unit/aiplatform/test_metadata_resources.py
@@ -37,6 +37,7 @@
     Execution as GapicExecution,
     Context as GapicContext,
     Artifact as GapicArtifact,
+    MetadataStore as GapicMetadataStore,
     AddContextArtifactsAndExecutionsResponse,
 )
 
@@ -361,6 +362,17 @@ def update_artifact_mock():
         yield update_artifact_mock
 
 
+@pytest.fixture
+def get_metadata_store_mock():
+    with patch.object(
+        MetadataServiceClient, "get_metadata_store"
+    ) as get_metadata_store_mock:
+        get_metadata_store_mock.return_value = GapicMetadataStore(
+            name=_TEST_METADATA_STORE
+        )
+        yield get_metadata_store_mock
+
+
 @pytest.mark.usefixtures("google_auth_mock")
 class TestContext:
     def setup_method(self):
@@ -799,7 +811,9 @@ def test_add_vertex_model(
             events=[Event(artifact=_TEST_ARTIFACT_NAME, type_=Event.Type.OUTPUT)],
         )
 
-    @pytest.mark.usefixtures("get_execution_mock", "get_model_with_version_mock")
+    @pytest.mark.usefixtures(
+        "get_execution_mock", "get_model_with_version_mock", "get_metadata_store_mock"
+    )
     def test_add_vertex_model_not_resolved(
         self, add_execution_events_mock, list_artifact_empty_mock, create_artifact_mock
     ):
