|
19 | 19 | import static com.google.common.truth.Truth.assertThat;
|
20 | 20 | import static junit.framework.TestCase.assertNotNull;
|
21 | 21 |
|
| 22 | +import com.google.cloud.aiplatform.v1.PipelineServiceClient; |
| 23 | +import com.google.cloud.aiplatform.v1.PipelineServiceSettings; |
| 24 | +import com.google.cloud.aiplatform.v1.TrainingPipeline; |
| 25 | +import com.google.cloud.aiplatform.v1.TrainingPipelineName; |
22 | 26 | import java.io.ByteArrayOutputStream;
|
23 | 27 | import java.io.IOException;
|
24 | 28 | import java.io.PrintStream;
|
@@ -72,7 +76,25 @@ public void tearDown()
|
72 | 76 | // Assert
|
73 | 77 | String cancelResponse = bout.toString();
|
74 | 78 | assertThat(cancelResponse).contains("Cancelled the Training Pipeline");
|
75 |
| - TimeUnit.MINUTES.sleep(2); |
| 79 | + |
| 80 | + PipelineServiceSettings pipelineServiceSettings = |
| 81 | + PipelineServiceSettings.newBuilder() |
| 82 | + .setEndpoint("us-central1-aiplatform.googleapis.com:443") |
| 83 | + .build(); |
| 84 | + |
| 85 | + try (PipelineServiceClient pipelineServiceClient = |
| 86 | + PipelineServiceClient.create(pipelineServiceSettings)) { |
| 87 | + String location = "us-central1"; |
| 88 | + TrainingPipelineName trainingPipelineName = |
| 89 | + TrainingPipelineName.of(PROJECT, location, trainingPipelineId); |
| 90 | + |
| 91 | + TrainingPipeline trainingPipelineResponse = |
| 92 | + pipelineServiceClient.getTrainingPipeline(trainingPipelineName); |
| 93 | + while (!trainingPipelineResponse.getState().name().contains("STATE_CANCELLED")) { |
| 94 | + TimeUnit.SECONDS.sleep(30); |
| 95 | + trainingPipelineResponse = pipelineServiceClient.getTrainingPipeline(trainingPipelineName); |
| 96 | + } |
| 97 | + } |
76 | 98 |
|
77 | 99 | // Delete the Training Pipeline
|
78 | 100 | DeleteTrainingPipelineSample.deleteTrainingPipelineSample(PROJECT, trainingPipelineId);
|
|
0 commit comments