8000 feat(tpu): add tpu queued resources list sample (#9614) · invertase/java-docs-samples@29ccf1c · GitHub
[go: up one dir, main page]

Skip to content

Commit 29ccf1c

Browse files
feat(tpu): add tpu queued resources list sample (GoogleCloudPlatform#9614)
* Changed package, added information to CODEOWNERS * Added information to CODEOWNERS * Added timeout * Fixed parameters for test * Fixed DeleteTpuVm and naming * Added comment, created Util class * Fixed naming * Fixed whitespace * Split PR into smaller, deleted redundant code * Implemented tpu_queued_resources_create, tpu_queued_resources_get, tpu_queued_resources_delete_force and tpu_queued_resources_delete samples, created tests * Implemented tpu_queued_resources_list sample, created test * Fixed test * Fixed tests, deleted cleanup method * Fixed test * Fixed imports
1 parent cfd36a0 commit 29ccf1c

File tree

3 files changed

+89
-1
lines changed

3 files changed

+89
-1
lines changed

tpu/src/main/java/tpu/DeleteForceQueuedResource.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -74,4 +74,4 @@ public static void deleteForceQueuedResource(
7474
}
7575
}
7676
}
77-
//[END tpu_queued_resources_delete_force]
77+
//[END tpu_queued_resources_delete_force]
Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,55 @@
1+
/*
2+
* Copyright 2024 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package tpu;
18+
19+
//[START tpu_queued_resources_list]
20+
import com.google.cloud.tpu.v2alpha1.ListQueuedResourcesRequest;
21+
import com.google.cloud.tpu.v2alpha1.QueuedResource;
22+
import com.google.cloud.tpu.v2alpha1.TpuClient;
23+
import com.google.cloud.tpu.v2alpha1.TpuClient.ListQueuedResourcesPage;
24+
import java.io.IOException;
25+
26+
public class ListQueuedResources {
27+
public static void main(String[] args) throws IOException {
28+
// TODO(developer): Replace these variables before running the sample.
29+
// Project ID or project number of the Google Cloud project.
30+
String projectId = "YOUR_PROJECT_ID";
31+
// The zone in which the TPU was created.
32+
String zone = "us-central1-a";
33+
34+
listQueuedResources(projectId, zone);
35+
}
36+
37+
// List Queued Resources.
38+
public static ListQueuedResourcesPage listQueuedResources(
39+
String projectId, String zone) throws IOException {
40+
String parent = String.format("projects/%s/locations/%s", projectId, zone);
41+
// Initialize client that will be used to send requests. This client only needs to be created
42+
// once, and can be reused for multiple requests.
43+
try (TpuClient tpuClient = TpuClient.create()) {
44+
ListQueuedResourcesRequest request =
45+
ListQueuedResourcesRequest.newBuilder().setParent(parent).build();
46+
ListQueuedResourcesPage response = tpuClient.listQueuedResources(request).getPage();
47+
48+
for (QueuedResource queuedResource : response.iterateAll()) {
49+
System.out.println(queuedResource.getName());
50+
}
51+
return response;
52+
}
53+
}
54+
}
55+
//[END tpu_queued_resources_list]

tpu/src/test/java/tpu/QueuedResourceIT.java

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616

1717
package tpu;
1818

19+
import static com.google.common.truth.Truth.assertThat;
1920
import static org.junit.Assert.assertEquals;
2021
import static org.mockito.ArgumentMatchers.anyLong;
2122
import static org.mockito.Mockito.any;
@@ -29,10 +30,15 @@
2930
import com.google.cloud.tpu.v2alpha1.CreateQueuedResourceRequest;
3031
import com.google.cloud.tpu.v2alpha1.DeleteQueuedResourceRequest;
3132
import com.google.cloud.tpu.v2alpha1.GetQueuedResourceRequest;
33+
import com.google.cloud.tpu.v2alpha1.ListQueuedResourcesRequest;
3234
import com.google.cloud.tpu.v2alpha1.QueuedResource;
3335
import com.google.cloud.tpu.v2alpha1.TpuClient;
36+
import com.google.cloud.tpu.v2alpha1.TpuClient.ListQueuedResourcesPage;
37+
import com.google.cloud.tpu.v2alpha1.TpuClient.ListQueuedResourcesPagedResponse;
3438
import com.google.cloud.tpu.v2alpha1.TpuSettings;
3539
import java.io.IOException;
40+
import java.util.Arrays;
41+
import java.util.List;
3642
import java.util.concurrent.ExecutionException;
3743
import java.util.concurrent.TimeUnit;
3844
import org.junit.jupiter.api.Test;
@@ -120,6 +126,33 @@ public void testGetQueuedResource() throws IOException {
120126
}
121127
}
122128

129+
@Test
130+
public void testListTpuVm() throws IOException {
131+
try (MockedStatic<TpuClient> mockedTpuClient = mockStatic(TpuClient.class)) {
132+
QueuedResource queuedResource1 = mock(QueuedResource.class);
133+
QueuedResource queuedResource2 = mock(QueuedResource.class);
134+
List<QueuedResource> mockListQueuedResources =
135+
Arrays.asList(queuedResource1, queuedResource2);
136+
137+
TpuClient mockClient = mock(TpuClient.class);
138+
mockedTpuClient.when(TpuClient::create).thenReturn(mockClient);
139+
ListQueuedResourcesPagedResponse mockListQueuedResourcesResponse =
140+
mock(ListQueuedResourcesPagedResponse.class);
141+
when(mockClient.listQueuedResources(any(ListQueuedResourcesRequest.class)))
142+
.thenReturn(mockListQueuedResourcesResponse);
143+
ListQueuedResourcesPage mockQueuedResourcesPage =
144+
mock(ListQueuedResourcesPage.class);
145+
when(mockListQueuedResourcesResponse.getPage()).thenReturn(mockQueuedResourcesPage);
146+
when(mockQueuedResourcesPage.getValues()).thenReturn(mockListQueuedResources);
147+
148+
ListQueuedResourcesPage returnedList =
149+
ListQueuedResources.listQueuedResources(PROJECT_ID, ZONE);
150+
151+
assertThat(returnedList.getValues()).isEqualTo(mockListQueuedResources);
152+
verify(mockClient, times(1)).listQueuedResources(any(ListQueuedResourcesRequest.class));
153+
}
154+
}
155+
123156
@Test
124157
public void testDeleteForceQueuedResource()
125158
throws IOException, InterruptedException, ExecutionException {

0 commit comments

Comments
 (0)
0