|
16 | 16 |
|
17 | 17 | package tpu;
|
18 | 18 |
|
| 19 | +import static com.google.common.truth.Truth.assertThat; |
19 | 20 | import static org.junit.Assert.assertEquals;
|
20 | 21 | import static org.mockito.ArgumentMatchers.anyLong;
|
21 | 22 | import static org.mockito.Mockito.any;
|
|
29 | 30 | import com.google.cloud.tpu.v2alpha1.CreateQueuedResourceRequest;
|
30 | 31 | import com.google.cloud.tpu.v2alpha1.DeleteQueuedResourceRequest;
|
31 | 32 | import com.google.cloud.tpu.v2alpha1.GetQueuedResourceRequest;
|
| 33 | +import com.google.cloud.tpu.v2alpha1.ListQueuedResourcesRequest; |
32 | 34 | import com.google.cloud.tpu.v2alpha1.QueuedResource;
|
33 | 35 | 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; |
34 | 38 | import com.google.cloud.tpu.v2alpha1.TpuSettings;
|
35 | 39 | import java.io.IOException;
|
| 40 | +import java.util.Arrays; |
| 41 | +import java.util.List; |
36 | 42 | import java.util.concurrent.ExecutionException;
|
37 | 43 | import java.util.concurrent.TimeUnit;
|
38 | 44 | import org.junit.jupiter.api.Test;
|
@@ -120,6 +126,33 @@ public void testGetQueuedResource() throws IOException {
|
120 | 126 | }
|
121 | 127 | }
|
122 | 128 |
|
| 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 | + |
123 | 156 | @Test
|
124 | 157 | public void testDeleteForceQueuedResource()
|
125 | 158 | throws IOException, InterruptedException, ExecutionException {
|
|
0 commit comments