feat: Add metadata SDK samples for list artifact and list execution (#1514)

SinaChavoshi · web-flow · commit c0d01f1d7c9d · 2022-07-21T17:18:34.000-07:00
* feat: Add metadata SDK samples for list artifact and execution

* update default execution display name

* Fix lint issues.

* update tests based on review feedback to return a list of more than one item

* simplify sample file names
diff --git a/samples/model-builder/conftest.py b/samples/model-builder/conftest.py
@@ -590,6 +590,15 @@ def mock_create_execution(mock_execution):
         yield mock_create_execution
 
 
+@pytest.fixture
+def mock_list_execution(mock_execution):
+    with patch.object(aiplatform.Execution, "list") as mock_list_execution:
+        # Returning list of 2 executions to avoid confusion with get method
+        # which returns one unique execution.
+        mock_list_execution.return_value = [mock_execution, mock_execution]
+        yield mock_list_execution
+
+
 @pytest.fixture
 def mock_get_artifact(mock_artifact):
     with patch.object(aiplatform, "Artifact") as mock_get_artifact:
@@ -625,6 +634,15 @@ def mock_create_artifact(mock_artifact):
         yield mock_create_artifact
 
 
+@pytest.fixture
+def mock_list_artifact(mock_artifact):
+    with patch.object(aiplatform.Artifact, "list") as mock_list_artifact:
+        # Returning list of 2 artifacts to avoid confusion with get method
+        # which returns one unique artifact.
+        mock_list_artifact.return_value = [mock_artifact, mock_artifact]
+        yield mock_list_artifact
+
+
 @pytest.fixture
 def mock_start_run(mock_experiment_run):
     with patch.object(aiplatform, "start_run") as mock_start_run:
diff --git a/samples/model-builder/experiment_tracking/list_artifact_sample.py b/samples/model-builder/experiment_tracking/list_artifact_sample.py
@@ -0,0 +1,36 @@
+# Copyright 2022 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Optional
+
+from google.cloud import aiplatform
+
+
+#  [START aiplatform_sdk_create_artifact_with_sdk_sample]
+def list_artifact_sample(
+    project: str,
+    location: str,
+    display_name_fitler: Optional[str] = "display_name=\"my_model_*\"",
+    create_date_filter:  Optional[str] = "create_time>\"2022-06-11T12:30:00-08:00\"",
+):
+    aiplatform.init(
+        project=project,
+        location=location)
+
+    combined_filters = f"{display_name_fitler} AND {create_date_filter}"
+
+    return aiplatform.Artifact.list(filter=combined_filters)
+
+
+#  [END aiplatform_sdk_create_artifact_with_sdk_sample]
diff --git a/samples/model-builder/experiment_tracking/list_artifact_sample_test.py b/samples/model-builder/experiment_tracking/list_artifact_sample_test.py
@@ -0,0 +1,33 @@
+# Copyright 2022 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import list_artifact_sample
+
+import test_constants as constants
+
+
+def test_list_artifact_with_sdk_sample(mock_artifact, mock_list_artifact):
+    artifacts = list_artifact_sample.list_artifact_sample(
+        project=constants.PROJECT,
+        location=constants.LOCATION,
+        display_name_fitler=constants.DISPLAY_NAME,
+        create_date_filter=constants.CREATE_DATE,
+    )
+
+    mock_list_artifact.assert_called_with(
+        filter=f"{constants.DISPLAY_NAME} AND {constants.CREATE_DATE}"
+    )
+    assert len(artifacts) == 2
+    assert artifacts[0] is mock_artifact
+    assert artifacts[1] is mock_artifact
diff --git a/samples/model-builder/experiment_tracking/list_execution_sample.py b/samples/model-builder/experiment_tracking/list_execution_sample.py
@@ -0,0 +1,36 @@
+# Copyright 2022 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Optional
+
+from google.cloud import aiplatform
+
+
+#  [START aiplatform_sdk_create_execution_with_sdk_sample]
+def list_execution_sample(
+    project: str,
+    location: str,
+    display_name_fitler: Optional[str] = "display_name=\"my_execution_*\"",
+    create_date_filter:  Optional[str] = "create_time>\"2022-06-11T12:30:00-08:00\"",
+):
+    aiplatform.init(
+        project=project,
+        location=location)
+
+    combined_filters = f"{display_name_fitler} AND {create_date_filter}"
+
+    return aiplatform.Execution.list(filter=combined_filters)
+
+
+#  [END aiplatform_sdk_create_execution_with_sdk_sample]
diff --git a/samples/model-builder/experiment_tracking/list_execution_sample_test.py b/samples/model-builder/experiment_tracking/list_execution_sample_test.py
@@ -0,0 +1,33 @@
+# Copyright 2022 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import list_execution_sample
+
+import test_constants as constants
+
+
+def test_list_execution_sample(mock_execution, mock_list_execution):
+    executions = list_execution_sample.list_execution_sample(
+        project=constants.PROJECT,
+        location=constants.LOCATION,
+        display_name_fitler=constants.DISPLAY_NAME,
+        create_date_filter=constants.CREATE_DATE,
+    )
+
+    mock_list_execution.assert_called_with(
+        filter=f"{constants.DISPLAY_NAME} AND {constants.CREATE_DATE}"
+    )
+    assert len(executions) == 2
+    assert executions[0] is mock_execution
+    assert executions[1] is mock_execution
diff --git a/samples/model-builder/test_constants.py b/samples/model-builder/test_constants.py
@@ -28,6 +28,8 @@
 DISPLAY_NAME = str(uuid4())  # Create random display name
 DISPLAY_NAME_2 = str(uuid4())
 
+CREATE_DATE = "2022-06-11T12:30:00-08:00"
+
 STAGING_BUCKET = "gs://my-staging-bucket"
 EXPERIMENT_NAME = "fraud-detection-trial-72"
 CREDENTIALS = credentials.AnonymousCredentials()
