docs: add a sample for get_experiment_run_artifacts

jaycee-li · copybara-github · commit 7266352ddddf · 2022-10-25T15:43:49.000-07:00
PiperOrigin-RevId: 483795137
diff --git a/samples/model-builder/conftest.py b/samples/model-builder/conftest.py
@@ -649,6 +649,12 @@ def mock_classification_metrics():
     yield mock
 
 
+@pytest.fixture
+def mock_artifacts():
+    mock = MagicMock()
+    yield mock
+
+
 @pytest.fixture
 def mock_get_execution(mock_execution):
     with patch.object(aiplatform, "Execution") as mock_get_execution:
@@ -903,6 +909,15 @@ def mock_get_classification_metrics(mock_classification_metrics, mock_experiment
         yield mock_get_classification_metrics
 
 
+@pytest.fixture
+def mock_get_artifacts(mock_artifacts, mock_experiment_run):
+    with patch.object(
+        mock_experiment_run, "get_artifacts"
+    ) as mock_get_artifacts:
+        mock_get_artifacts.return_value = mock_artifacts
+        yield mock_get_artifacts
+
+
 """
 ----------------------------------------------------------------------------
 Model Versioning Fixtures
diff --git a/samples/model-builder/experiment_tracking/get_experiment_run_artifacts_sample.py b/samples/model-builder/experiment_tracking/get_experiment_run_artifacts_sample.py
@@ -0,0 +1,37 @@
+# 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 List, Union
+
+from google.cloud import aiplatform
+from google.cloud.aiplatform.metadata import artifact
+
+
+#  [START aiplatform_sdk_get_experiment_run_artifacts_sample]
+def get_experiment_run_artifacts_sample(
+    run_name: str,
+    experiment: Union[str, aiplatform.Experiment],
+    project: str,
+    location: str,
+) -> List[artifact.Artifact]:
+    experiment_run = aiplatform.ExperimentRun(
+        run_name=run_name,
+        experiment=experiment,
+        project=project,
+        location=location,
+    )
+
+    return experiment_run.get_artifacts()
+
+#  [END aiplatform_sdk_get_experiment_run_artifacts_sample]
diff --git a/samples/model-builder/experiment_tracking/get_experiment_run_artifacts_sample_test.py b/samples/model-builder/experiment_tracking/get_experiment_run_artifacts_sample_test.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.
+
+import get_experiment_run_artifacts_sample
+
+import pytest
+
+import test_constants as constants
+
+
+@pytest.mark.usefixtures("mock_get_run")
+def test_get_experiment_run_artifact_sample(mock_get_artifacts, mock_artifacts):
+
+    artifacts = (
+        get_experiment_run_artifacts_sample.get_experiment_run_artifacts_sample(
+            run_name=constants.EXPERIMENT_RUN_NAME,
+            experiment=constants.EXPERIMENT_NAME,
+            project=constants.PROJECT,
+            location=constants.LOCATION,
+        )
+    )
+
+    mock_get_artifacts.assert_called_with()
+
+    assert artifacts is mock_artifacts
