docs: Fix incorrect field in log_time_series_metrics_sample

PiperOrigin-RevId: 605759888
googleapis · Feb 10, 2024 · f249353 · f249353
1 parent 90d95d7
commit f249353
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 5 deletions.
diff --git a/samples/model-builder/experiment_tracking/log_time_series_metrics_sample.py b/samples/model-builder/experiment_tracking/log_time_series_metrics_sample.py
@@ -30,7 +30,7 @@ def log_time_series_metrics_sample(
 ):
     aiplatform.init(experiment=experiment_name, project=project, location=location)
 
-    aiplatform.start_run(run_name=run_name, resume=True)
+    aiplatform.start_run(run=run_name, resume=True)
 
     aiplatform.log_time_series_metrics(metrics=metrics, step=step, wall_time=wall_time)
 

diff --git a/samples/model-builder/experiment_tracking/log_time_series_metrics_sample_test.py b/samples/model-builder/experiment_tracking/log_time_series_metrics_sample_test.py
@@ -12,14 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import pytest
-
 from experiment_tracking import log_time_series_metrics_sample
 import test_constants as constants
 
 
-@pytest.mark.usefixtures("mock_sdk_init", "mock_start_run")
-def test_log_time_series_metrics_sample(mock_log_time_series_metrics):
+def test_log_time_series_metrics_sample(
+        mock_log_time_series_metrics, mock_start_run, mock_sdk_init):
 
     log_time_series_metrics_sample.log_time_series_metrics_sample(
         experiment_name=constants.EXPERIMENT_NAME,
@@ -31,6 +29,17 @@ def test_log_time_series_metrics_sample(mock_log_time_series_metrics):
         location=constants.LOCATION,
     )
 
+    mock_sdk_init.assert_called_with(
+        experiment=constants.EXPERIMENT_NAME,
+        project=constants.PROJECT,
+        location=constants.LOCATION
+    )
+
+    mock_start_run.assert_called_with(
+        run=constants.EXPERIMENT_RUN_NAME,
+        resume=True,
+    )
+
     mock_log_time_series_metrics.assert_called_with(
         metrics=constants.METRICS,
         step=constants.STEP,