1 /*
2  * Copyright 2020 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  *     https://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 com.google.cloud.notebooks.v1beta1.it;
18 
19 import static org.junit.Assert.assertEquals;
20 import static org.junit.Assert.assertTrue;
21 
22 import com.google.api.FieldBehavior;
23 import com.google.cloud.ServiceOptions;
24 import com.google.cloud.notebooks.v1beta1.ContainerImage;
25 import com.google.cloud.notebooks.v1beta1.CreateEnvironmentRequest;
26 import com.google.cloud.notebooks.v1beta1.CreateInstanceRequest;
27 import com.google.cloud.notebooks.v1beta1.DeleteEnvironmentRequest;
28 import com.google.cloud.notebooks.v1beta1.DeleteInstanceRequest;
29 import com.google.cloud.notebooks.v1beta1.Environment;
30 import com.google.cloud.notebooks.v1beta1.GetEnvironmentRequest;
31 import com.google.cloud.notebooks.v1beta1.GetInstanceRequest;
32 import com.google.cloud.notebooks.v1beta1.Instance;
33 import com.google.cloud.notebooks.v1beta1.ListEnvironmentsRequest;
34 import com.google.cloud.notebooks.v1beta1.ListInstancesRequest;
35 import com.google.cloud.notebooks.v1beta1.NotebookServiceClient;
36 import com.google.cloud.notebooks.v1beta1.ResetInstanceRequest;
37 import com.google.cloud.notebooks.v1beta1.SetInstanceAcceleratorRequest;
38 import com.google.cloud.notebooks.v1beta1.SetInstanceLabelsRequest;
39 import com.google.cloud.notebooks.v1beta1.SetInstanceMachineTypeRequest;
40 import com.google.cloud.notebooks.v1beta1.StartInstanceRequest;
41 import com.google.cloud.notebooks.v1beta1.StopInstanceRequest;
42 import java.io.IOException;
43 import java.util.UUID;
44 import java.util.concurrent.ExecutionException;
45 import org.junit.AfterClass;
46 import org.junit.BeforeClass;
47 import org.junit.Test;
48 
49 public class ITNotebookServiceClientTest {
50 
51   private static final String PROJECT_ID = ServiceOptions.getDefaultProjectId();
52   private static final String LOCATION = "us-central1-a";
53   private static final String PARENT = "projects/" + PROJECT_ID + "/locations/" + LOCATION;
54   private static NotebookServiceClient client;
55   private static final String ID = UUID.randomUUID().toString().substring(0, 8);
56   private static final String NOTEBOOK_PREFIX = "it-test-notebook";
57   private static final String NOTEBOOK_INSTANCE_ID = NOTEBOOK_PREFIX + "-instance-id-" + ID;
58   private static final String ENVIRONMENT_ID = NOTEBOOK_PREFIX + "-environment-id-" + ID;
59   private static final String INSTANCE_NAME = PARENT + "/instances/" + NOTEBOOK_INSTANCE_ID;
60   private static final String ENVIRONMENT_NAME = PARENT + "/environments/" + ENVIRONMENT_ID;
61   private static Instance expectedNotebookInstance;
62   private static Environment expectedEnvironmentResponse;
63   private static final String MACHINE_TYPE_A = "n1-standard-4";
64   private static final String MACHINE_TYPE_B = "n1-standard-2";
65 
66   @BeforeClass
setUp()67   public static void setUp() throws IOException, ExecutionException, InterruptedException {
68     // Create Test Notebook Instance
69     client = NotebookServiceClient.create();
70     Util.cleanUpNotebookInstances(client, PARENT, NOTEBOOK_PREFIX);
71 
72     ContainerImage containerImage =
73         ContainerImage.newBuilder().setRepository(FieldBehavior.OPTIONAL.name()).build();
74 
75     Environment environment =
76         Environment.newBuilder()
77             .setName(ENVIRONMENT_NAME)
78             .setContainerImage(containerImage)
79             .build();
80 
81     CreateEnvironmentRequest environmentRequest =
82         CreateEnvironmentRequest.newBuilder()
83             .setParent(PARENT)
84             .setEnvironmentId(ENVIRONMENT_ID)
85             .setEnvironment(environment)
86             .build();
87 
88     expectedEnvironmentResponse = client.createEnvironmentAsync(environmentRequest).get();
89 
90     Instance notebookInstance =
91         Instance.newBuilder()
92             .setContainerImage(containerImage)
93             .setMachineType(MACHINE_TYPE_A)
94             .build();
95     CreateInstanceRequest instanceRequest =
96         CreateInstanceRequest.newBuilder()
97             .setParent(PARENT)
98             .setInstanceId(NOTEBOOK_INSTANCE_ID)
99             .setInstance(notebookInstance)
100             .build();
101     expectedNotebookInstance = client.createInstanceAsync(instanceRequest).get();
102   }
103 
104   @AfterClass
tearDown()105   public static void tearDown() throws ExecutionException, InterruptedException {
106     // Delete Test Environment Instance
107     DeleteEnvironmentRequest deleteEnvironmentRequest =
108         DeleteEnvironmentRequest.newBuilder().setName(ENVIRONMENT_NAME).build();
109     client.deleteEnvironmentAsync(deleteEnvironmentRequest).get();
110     // Delete Test Notebook Instance
111     DeleteInstanceRequest deleteInstanceRequest =
112         DeleteInstanceRequest.newBuilder().setName(INSTANCE_NAME).build();
113     client.deleteInstanceAsync(deleteInstanceRequest).get();
114     client.close();
115   }
116 
stopInstance()117   private String stopInstance() throws ExecutionException, InterruptedException {
118     StopInstanceRequest request = StopInstanceRequest.newBuilder().setName(INSTANCE_NAME).build();
119     Instance response = client.stopInstanceAsync(request).get();
120     return response.getState().name();
121   }
122 
startInstance()123   private String startInstance() throws ExecutionException, InterruptedException {
124     StartInstanceRequest request = StartInstanceRequest.newBuilder().setName(INSTANCE_NAME).build();
125     Instance response = client.startInstanceAsync(request).get();
126     return response.getState().name();
127   }
128 
129   @Test
testGetInstance()130   public void testGetInstance() {
131     GetInstanceRequest request = GetInstanceRequest.newBuilder().setName(INSTANCE_NAME).build();
132     Instance actualNotebookInstance = client.getInstance(request);
133     assertEquals(
134         expectedNotebookInstance.getContainerImage(), actualNotebookInstance.getContainerImage());
135     assertEquals(expectedNotebookInstance.getName(), actualNotebookInstance.getName());
136     assertEquals(expectedNotebookInstance.getNetwork(), actualNotebookInstance.getNetwork());
137     assertEquals(expectedNotebookInstance.getSubnet(), actualNotebookInstance.getSubnet());
138   }
139 
140   @Test
testListInstances()141   public void testListInstances() {
142     ListInstancesRequest request = ListInstancesRequest.newBuilder().setParent(PARENT).build();
143     for (Instance element : client.listInstances(request).iterateAll()) {
144       if (element.getName().equals(NOTEBOOK_INSTANCE_ID)) {
145         assertEquals(expectedNotebookInstance.getContainerImage(), element.getContainerImage());
146         assertEquals(expectedNotebookInstance.getName(), element.getName());
147         assertEquals(expectedNotebookInstance.getMachineType(), element.getMachineType());
148         assertEquals(expectedNotebookInstance.getNetwork(), element.getNetwork());
149         assertEquals(expectedNotebookInstance.getSubnet(), element.getSubnet());
150       }
151     }
152   }
153 
154   @Test
testGetEnvironment()155   public void testGetEnvironment() {
156     GetEnvironmentRequest environmentRequest =
157         GetEnvironmentRequest.newBuilder().setName(ENVIRONMENT_NAME).build();
158     Environment expectedEnvironmentResponse = client.getEnvironment(environmentRequest);
159     assertEquals(expectedEnvironmentResponse.getName(), expectedEnvironmentResponse.getName());
160     assertEquals(
161         expectedEnvironmentResponse.getContainerImage(),
162         expectedEnvironmentResponse.getContainerImage());
163   }
164 
165   @Test
testListEnvironment()166   public void testListEnvironment() {
167     ListEnvironmentsRequest request =
168         ListEnvironmentsRequest.newBuilder().setParent(PARENT).build();
169     for (Environment element : client.listEnvironments(request).iterateAll()) {
170       if (element.getName().equals(ENVIRONMENT_ID)) {
171         assertEquals(expectedEnvironmentResponse.getName(), element.getName());
172         assertEquals(expectedEnvironmentResponse.getContainerImage(), element.getContainerImage());
173       }
174     }
175   }
176 
177   @Test
testSetInstanceLabels()178   public void testSetInstanceLabels() throws ExecutionException, InterruptedException {
179     SetInstanceLabelsRequest setInstanceLabelsRequest =
180         SetInstanceLabelsRequest.newBuilder().setName(INSTANCE_NAME).putLabels("a", "100").build();
181     Instance response = client.setInstanceLabelsAsync(setInstanceLabelsRequest).get();
182     assertEquals(response.getLabelsMap().get("a"), "100");
183   }
184 
185   @Test
testSetInstanceMachineType()186   public void testSetInstanceMachineType() throws ExecutionException, InterruptedException {
187     this.stopInstance();
188     SetInstanceMachineTypeRequest request =
189         SetInstanceMachineTypeRequest.newBuilder()
190             .setName(INSTANCE_NAME)
191             .setMachineType(MACHINE_TYPE_B)
192             .build();
193     Instance response = client.setInstanceMachineTypeAsync(request).get();
194     assertTrue(response.getMachineType().endsWith(MACHINE_TYPE_B));
195     this.startInstance();
196   }
197 
198   @Test
testStartInstance()199   public void testStartInstance() throws ExecutionException, InterruptedException {
200     String state = this.stopInstance();
201     state = this.startInstance();
202     assertEquals(state, "PROVISIONING");
203   }
204 
205   @Test
testStopInstance()206   public void testStopInstance() throws ExecutionException, InterruptedException {
207     String state = this.stopInstance();
208     assertEquals(state, "STOPPED");
209     state = this.startInstance();
210   }
211 
212   @Test
testSetInstanceAccelarator()213   public void testSetInstanceAccelarator() throws ExecutionException, InterruptedException {
214     this.stopInstance();
215     SetInstanceAcceleratorRequest request =
216         SetInstanceAcceleratorRequest.newBuilder()
217             .setName(INSTANCE_NAME)
218             .setType(Instance.AcceleratorType.NVIDIA_TESLA_P4)
219             .setCoreCount(1L)
220             .build();
221     Instance response = client.setInstanceAcceleratorAsync(request).get();
222     this.startInstance();
223     assertEquals(
224         response.getAcceleratorConfig().getType().name(),
225         Instance.AcceleratorType.NVIDIA_TESLA_P4.name());
226   }
227 
228   @Test
testResetInstance()229   public void testResetInstance() throws ExecutionException, InterruptedException {
230     ResetInstanceRequest request = ResetInstanceRequest.newBuilder().setName(INSTANCE_NAME).build();
231     Instance response = client.resetInstanceAsync(request).get();
232     assertTrue(response.getMachineType().endsWith(MACHINE_TYPE_B));
233   }
234 }
235