chore: add blocking call to forecasting test

sararob · copybara-github · commit f837e0e49f2b · 2023-04-12T15:31:49.000-07:00
PiperOrigin-RevId: 523822564
diff --git a/tests/system/aiplatform/test_e2e_forecasting.py b/tests/system/aiplatform/test_e2e_forecasting.py
@@ -120,6 +120,7 @@ def test_end_to_end_forecasting(self, shared_state, training_job):
             # resources.append(batch_prediction_job)
 
             # batch_prediction_job.wait()
+            model.wait()
             assert job.state == pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED
             # assert batch_prediction_job.state == job_state.JobState.JOB_STATE_SUCCEEDED
         finally:
