chore: Add updateSchedule test to PipelineJobSchedule system test.

vertex-sdk-bot · copybara-github · commit 1b91e9f80f59 · 2023-06-20T14:42:52.000-07:00
PiperOrigin-RevId: 542053580
diff --git a/tests/system/aiplatform/test_pipeline_job_schedule.py b/tests/system/aiplatform/test_pipeline_job_schedule.py
@@ -37,7 +37,7 @@
 class TestPreviewPipelineJobSchedule(e2e_base.TestEndToEnd):
     _temp_prefix = "tmpvrtxsdk-e2e-pjs"
 
-    def test_create_get_list(self, shared_state):
+    def test_create_get_pause_resume_update_list(self, shared_state):
         # Components:
         def train(
             number_of_epochs: int,
@@ -77,22 +77,34 @@ def training_pipeline(number_of_epochs: int = 2):
         )
 
         max_run_count = 2
+        cron_expression = "*/5 * * * *"
         pipeline_job_schedule.create(
-            cron_expression="*/5 * * * *",
+            cron_expression=cron_expression,
             max_run_count=max_run_count,
             max_concurrent_run_count=2,
         )
 
         shared_state.setdefault("resources", []).append(pipeline_job_schedule)
 
+        # Pausing the pipeline job schedule.
         pipeline_job_schedule.pause()
         assert pipeline_job_schedule.state == gca_schedule.Schedule.State.PAUSED
 
+        # Before updating, confirm the cron_expression is correctly set from the create step.
+        assert pipeline_job_schedule.cron_expression == cron_expression
+
+        # Updating the pipeline job schedule.
+        new_cron_expression = "* * * * *"
+        pipeline_job_schedule.update(cron_expression=new_cron_expression)
+        assert pipeline_job_schedule.cron_expression == new_cron_expression
+
+        # Resuming the pipeline job schedule.
         pipeline_job_schedule.resume(catch_up=True)
         assert pipeline_job_schedule.state == gca_schedule.Schedule.State.ACTIVE
 
         pipeline_job_schedule.wait()
 
+        # Confirming that correct number of runs were scheduled and completed by this pipeline job schedule.
         list_jobs_with_read_mask = pipeline_job_schedule.list_jobs(
             enable_simple_view=True
         )
