Skip to content

Commit ed08095

Browse files
Sita04Sita04
committed
update snippets to wait for job completion
1 parent 695e875 commit ed08095

File tree

4 files changed

+68
-10
lines changed

4 files changed

+68
-10
lines changed

‎batch/snippets/src/test/java/BatchBasicIT.java

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -68,15 +68,18 @@ public static void setUp()
6868
.contains(
6969
"Successfully created the job: "
7070
+ String.format(
71-
"projects/%s/locations/%s/jobs/%s", PROJECT_ID, REGION, CONTAINER_JOB_NAME));
71+
"projects/%s/locations/%s/jobs/%s", PROJECT_ID, REGION, CONTAINER_JOB_NAME));
7272
CreateWithScriptNoMounting.createScriptJob(PROJECT_ID, REGION, SCRIPT_JOB_NAME);
7373
assertThat(stdOut.toString())
7474
.contains(
7575
"Successfully created the job: "
7676
+ String.format(
77-
"projects/%s/locations/%s/jobs/%s", PROJECT_ID, REGION, SCRIPT_JOB_NAME));
77+
"projects/%s/locations/%s/jobs/%s", PROJECT_ID, REGION, SCRIPT_JOB_NAME));
7878
TimeUnit.SECONDS.sleep(10);
7979

80+
Util.waitForJobCompletion(Util.getJob(CONTAINER_JOB_NAME, PROJECT_ID, REGION));
81+
Util.waitForJobCompletion(Util.getJob(SCRIPT_JOB_NAME, PROJECT_ID, REGION));
82+
8083
stdOut.close();
8184
System.setOut(out);
8285
}

‎batch/snippets/src/test/java/BatchBucketIT.java

Lines changed: 26 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -15,16 +15,17 @@
1515
import static com.google.common.truth.Truth.assertThat;
1616
import static com.google.common.truth.Truth.assertWithMessage;
1717

18+
import com.google.cloud.batch.v1.Job;
1819
import com.google.cloud.storage.Blob;
1920
import com.google.cloud.storage.Bucket;
2021
import com.google.cloud.storage.BucketInfo;
2122
import com.google.cloud.storage.Storage;
23+
import com.google.cloud.storage.StorageClass;
2224
import com.google.cloud.storage.StorageOptions;
2325
import java.io.ByteArrayOutputStream;
2426
import java.io.IOException;
2527
import java.io.PrintStream;
26-
import java.util.Arrays;
27-
import java.util.Base64;
28+
import java.nio.charset.StandardCharsets;
2829
import java.util.UUID;
2930
import java.util.concurrent.ExecutionException;
3031
import java.util.concurrent.TimeUnit;
@@ -44,7 +45,6 @@ public class BatchBucketIT {
4445
private static final String REGION = "europe-north1";
4546
private static String SCRIPT_JOB_NAME;
4647
private static String BUCKET_NAME;
47-
4848
private ByteArrayOutputStream stdOut;
4949

5050
// Check if the required environment variables are set.
@@ -81,15 +81,31 @@ public static void cleanup()
8181
ByteArrayOutputStream stdOut = new ByteArrayOutputStream();
8282
System.setOut(new PrintStream(stdOut));
8383

84+
// Delete bucket.
85+
Storage storage = StorageOptions.newBuilder().setProjectId(PROJECT_ID).build().getService();
86+
Bucket bucket = storage.get(BUCKET_NAME);
87+
for (Blob blob : storage.list(bucket.getName()).iterateAll()) {
88+
storage.delete(blob.getBlobId());
89+
}
90+
storage.delete(bucket.getName());
91+
System.out.println("Bucket " + bucket.getName() + " was deleted");
92+
93+
// Delete job.
8494
DeleteJob.deleteJob(PROJECT_ID, REGION, SCRIPT_JOB_NAME);
8595

8696
stdOut.close();
8797
System.setOut(out);
8898
}
8999

90100
private static void createBucket(String bucketName) {
91-
Storage storage = StorageOptions.getDefaultInstance().getService();
92-
storage.create(BucketInfo.of(bucketName));
101+
Storage storage = StorageOptions.newBuilder().setProjectId(PROJECT_ID).build().getService();
102+
StorageClass storageClass = StorageClass.COLDLINE;
103+
String location = "US";
104+
storage.create(
105+
BucketInfo.newBuilder(bucketName)
106+
.setStorageClass(storageClass)
107+
.setLocation(location)
108+
.build());
93109
}
94110

95111
@Before
@@ -109,19 +125,22 @@ public void testBucketJob()
109125
throws IOException, ExecutionException, InterruptedException, TimeoutException {
110126
CreateWithMountedBucket.createScriptJobWithBucket(PROJECT_ID, REGION, SCRIPT_JOB_NAME,
111127
BUCKET_NAME);
128+
Job job = Util.getJob(PROJECT_ID, REGION, SCRIPT_JOB_NAME);
129+
Util.waitForJobCompletion(job);
112130
assertThat(stdOut.toString()).contains("Successfully created the job");
113131
testBucketContent();
114132
}
115133

116134
public void testBucketContent() {
117135
String fileNameTemplate = "output_task_%s.txt";
118-
String fileContentTemplate = "Hello world from task {task_number}.\n";
136+
String fileContentTemplate;
119137

120138
Storage storage = StorageOptions.getDefaultInstance().getService();
121139
Bucket bucket = storage.get(BUCKET_NAME);
122140
for (int i = 0; i < 4; i++) {
141+
fileContentTemplate = String.format("Hello world from task %s.\n", i);
123142
Blob blob = bucket.get(String.format(fileNameTemplate, i));
124-
String content = Arrays.toString(Base64.getDecoder().decode(blob.getContent()));
143+
String content = new String(blob.getContent(), StandardCharsets.UTF_8);
125144
assertThat(fileContentTemplate).matches(content);
126145
}
127146
}

‎batch/snippets/src/test/java/BatchTemplateIT.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -218,6 +218,6 @@ public void testCreateWithTemplate()
218218
CreateWithTemplate.createWithTemplate(PROJECT_ID, REGION, SCRIPT_JOB_NAME,
219219
INSTANCE_TEMPLATE.getSelfLink());
220220
assertThat(stdOut.toString()).contains("Successfully created the job: ");
221+
Util.waitForJobCompletion(Util.getJob(SCRIPT_JOB_NAME, PROJECT_ID, REGION));
221222
}
222-
223223
}

‎batch/snippets/src/test/java/Util.java

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,12 @@
1212
// See the License for the specific language governing permissions and
1313
// limitations under the License.
1414

15+
import static com.google.common.truth.Truth.assertThat;
16+
17+
import com.google.cloud.batch.v1.BatchServiceClient;
18+
import com.google.cloud.batch.v1.Job;
19+
import com.google.cloud.batch.v1.JobName;
20+
import com.google.cloud.batch.v1.JobStatus.State;
1521
import com.google.cloud.compute.v1.DeleteInstanceTemplateRequest;
1622
import com.google.cloud.compute.v1.InstanceTemplate;
1723
import com.google.cloud.compute.v1.InstanceTemplatesClient;
@@ -21,13 +27,16 @@
2127
import java.time.Instant;
2228
import java.time.OffsetDateTime;
2329
import java.time.temporal.ChronoUnit;
30+
import java.util.List;
2431
import java.util.concurrent.ExecutionException;
2532
import java.util.concurrent.TimeUnit;
2633
import java.util.concurrent.TimeoutException;
2734

2835
public class Util {
2936

3037
private static final int DELETION_THRESHOLD_TIME_HOURS = 24;
38+
private static final List<State> WAIT_STATES = List.of(State.STATE_UNSPECIFIED, State.QUEUED,
39+
State.RUNNING, State.SCHEDULED);
3140

3241
// Delete templates which starts with the given prefixToDelete and
3342
// has creation timestamp >24 hours.
@@ -78,4 +87,31 @@ private static void deleteInstanceTemplate(String projectId, String templateName
7887
.get(3, TimeUnit.MINUTES);
7988
}
8089
}
90+
91+
public static Job getJob(String projectId, String region, String jobName) throws IOException {
92+
try (BatchServiceClient batchServiceClient = BatchServiceClient.create()) {
93+
return
94+
batchServiceClient.getJob(
95+
JobName.newBuilder()
96+
.setProject(projectId)
97+
.setLocation(region)
98+
.setJob(jobName)
99+
.build());
100+
}
101+
}
102+
103+
public static void waitForJobCompletion(Job job)
104+
throws IOException, InterruptedException {
105+
String[] jobName = job.getName().split("/");
106+
Instant startTime = Instant.now();
107+
while (WAIT_STATES.contains(job.getStatus().getState())) {
108+
if (Instant.now().getEpochSecond() - startTime.getEpochSecond() > 300) {
109+
throw new Error("Timed out waiting for operation to complete.");
110+
}
111+
job = getJob(jobName[1], jobName[3], jobName[5]);
112+
TimeUnit.SECONDS.sleep(5);
113+
}
114+
job = getJob(jobName[1], jobName[3], job.getName().split("/")[5]);
115+
assertThat(job.getStatus().getState() == State.SUCCEEDED);
116+
}
81117
}

0 commit comments

Comments
 (0)