Use mocked TPUClient in tests

GoogleCloudPlatform · Dec 3, 2024 · 5875a3c · 5875a3c
1 parent 5eeb1f1
commit 5875a3c
Show file tree

Hide file tree

Showing 7 changed files with 252 additions and 172 deletions.
diff --git a/tpu/queuedResources/createQueuedResource.js b/tpu/queuedResources/createQueuedResource.js
@@ -16,24 +16,20 @@
 
 'use strict';
 
-async function main(
-  nodeName,
-  queuedResourceName,
-  zone,
-  tpuType,
-  tpuSoftwareVersion
-) {
+async function main(tpuClient) {
   // [START tpu_queued_resources_create]
-  // Import the TPU library
-  const {TpuClient} = require('@google-cloud/tpu').v2alpha1;
+  // Import the TPUClient
+  // TODO(developer): Uncomment below line before running the sample.
+  // const {TpuClient} = require('@google-cloud/tpu').v2alpha1;
   const {Node, NetworkConfig, QueuedResource} =
     require('@google-cloud/tpu').protos.google.cloud.tpu.v2alpha1;
 
   // Instantiate a tpuClient
-  const tpuClient = new TpuClient();
+  // TODO(developer): Uncomment below line before running the sample.
+  // tpuClient = new TpuClient();
 
   /**
-   * TODO(developer): Update/uncomment these variables before running the sample.
+   * TODO(developer): Update these variables before running the sample.
    */
   // Project ID or project number of the Google Cloud project, where you want to create queued resource.
   const projectId = await tpuClient.getProjectId();
@@ -45,24 +41,24 @@ async function main(
   const region = 'europe-west4';
 
   // The name for your queued resource.
-  // queuedResourceName = 'queued-resource-1';
+  const queuedResourceName = 'queued-resource-1';
 
   // The name for your node.
-  // nodeName = 'node-name-1';
+  const nodeName = 'node-name-1';
 
   // The zone in which to create the node.
   // For more information about supported TPU types for specific zones,
   // see https://cloud.google.com/tpu/docs/regions-zones
-  // zone = 'europe-west4-a';
+  const zone = 'europe-west4-a';
 
   // The accelerator type that specifies the version and size of the node you want to create.
   // For more information about supported accelerator types for each TPU version,
   // see https://cloud.google.com/tpu/docs/system-architecture-tpu-vm#versions.
-  // tpuType = 'v2-8';
+  const tpuType = 'v2-8';
 
   // Software version that specifies the version of the node runtime to install. For more information,
   // see https://cloud.google.com/tpu/docs/runtimes
-  // tpuSoftwareVersion = 'tpu-vm-tf-2.14.1';
+  const tpuSoftwareVersion = 'tpu-vm-tf-2.14.1';
 
   async function callCreateQueuedResource() {
     // Create a node
@@ -108,17 +104,21 @@ async function main(
     const [operation] = await tpuClient.createQueuedResource(request);
 
     // Wait for the create operation to complete.
-    await operation.promise();
+    const [response] = await operation.promise();
 
     // You can wait until TPU Node is READY,
-    // and check its status using getTpuVm() from `tpu_vm_get` sample.
+    // and check its status using callGetTpuVm() from `tpu_vm_get` sample.
     console.log(`Queued resource ${queuedResourceName} created.`);
+    return response;
   }
-  await callCreateQueuedResource();
+  return await callCreateQueuedResource();
   // [END tpu_queued_resources_create]
 }
 
-main(...process.argv.slice(2)).catch(err => {
-  console.error(err);
-  process.exitCode = 1;
-});
+module.exports = main;
+
+// TODO(developer): Uncomment below lines before running the sample.
+// main(...process.argv.slice(2)).catch(err => {
+//   console.error(err);
+//   process.exitCode = 1;
+// });
diff --git a/tpu/queuedResources/deleteQueuedResource.js b/tpu/queuedResources/deleteQueuedResource.js
@@ -16,25 +16,27 @@
 
 'use strict';
 
-async function main(queuedResourceName, zone) {
+async function main(tpuClient) {
   // [START tpu_queued_resources_delete]
-  // Import the TPU library
-  const {TpuClient} = require('@google-cloud/tpu').v2alpha1;
+  // Import the TPUClient
+  // TODO(developer): Uncomment below line before running the sample.
+  // const {TpuClient} = require('@google-cloud/tpu').v2alpha1;
 
   // Instantiate a tpuClient
-  const tpuClient = new TpuClient();
+  // TODO(developer): Uncomment below line before running the sample.
+  // tpuClient = new TpuClient();
 
   /**
-   * TODO(developer): Update/uncomment these variables before running the sample.
+   * TODO(developer): Update these variables before running the sample.
    */
   // Project ID or project number of the Google Cloud project, where you want to delete node.
   const projectId = await tpuClient.getProjectId();
 
   // The name of queued resource.
-  // queuedResourceName = 'queued-resource-1';
+  const queuedResourceName = 'queued-resource-1';
 
   // The zone of your queued resource.
-  // zone = 'europe-west4-a';
+  const zone = 'europe-west4-a';
 
   async function callDeleteTpuVM(nodeName) {
     const request = {
@@ -55,24 +57,28 @@ async function main(queuedResourceName, zone) {
     };
 
     // Retrive node name
-    const [response] = await tpuClient.getQueuedResource(request);
-    const nodeName = response.tpu.nodeSpec[0].nodeId;
+    const [queuedResource] = await tpuClient.getQueuedResource(request);
+    const nodeName = queuedResource.tpu.nodeSpec[0].nodeId;
 
     // Before deleting the queued resource it is required to delete the TPU VM.
     await callDeleteTpuVM(nodeName);
 
     const [operation] = await tpuClient.deleteQueuedResource(request);
 
     // Wait for the delete operation to complete.
-    await operation.promise();
+    const [response] = await operation.promise();
 
     console.log(`Queued resource ${queuedResourceName} deleted.`);
+    return response;
   }
-  await callDeleteQueuedResource();
+  return await callDeleteQueuedResource();
   // [END tpu_queued_resources_delete]
 }
 
-main(...process.argv.slice(2)).catch(err => {
-  console.error(err);
-  process.exitCode = 1;
-});
+module.exports = main;
+
+// TODO(developer): Uncomment below lines before running the sample.
+// main(...process.argv.slice(2)).catch(err => {
+//   console.error(err);
+//   process.exitCode = 1;
+// });
diff --git a/tpu/queuedResources/forceDeleteQueuedResource.js b/tpu/queuedResources/forceDeleteQueuedResource.js
@@ -16,25 +16,27 @@
 
 'use strict';
 
-async function main(queuedResourceName, zone) {
+async function main(tpuClient) {
   // [START tpu_queued_resources_delete_force]
-  // Import the TPU library
-  const {TpuClient} = require('@google-cloud/tpu').v2alpha1;
+  // Import the TPUClient
+  // TODO(developer): Uncomment below line before running the sample.
+  // const {TpuClient} = require('@google-cloud/tpu').v2alpha1;
 
   // Instantiate a tpuClient
-  const tpuClient = new TpuClient();
+  // TODO(developer): Uncomment below line before running the sample.
+  // tpuClient = new TpuClient();
 
   /**
-   * TODO(developer): Update/uncomment these variables before running the sample.
+   * TODO(developer): Update these variables before running the sample.
    */
   // Project ID or project number of the Google Cloud project, where you want to delete node.
   const projectId = await tpuClient.getProjectId();
 
   // The name of queued resource.
-  // queuedResourceName = 'queued-resource-1';
+  const queuedResourceName = 'queued-resource-1';
 
   // The zone of your queued resource.
-  // zone = 'europe-west4-a';
+  const zone = 'europe-west4-a';
 
   async function callForceDeleteQueuedResource() {
     const request = {
@@ -45,15 +47,19 @@ async function main(queuedResourceName, zone) {
     const [operation] = await tpuClient.deleteQueuedResource(request);
 
     // Wait for the delete operation to complete.
-    await operation.promise();
+    const [response] = await operation.promise();
 
     console.log(`Queued resource ${queuedResourceName} deletion forced.`);
+    return response;
   }
-  await callForceDeleteQueuedResource();
+  return await callForceDeleteQueuedResource();
   // [END tpu_queued_resources_delete_force]
 }
 
-main(...process.argv.slice(2)).catch(err => {
-  console.error(err);
-  process.exitCode = 1;
-});
+module.exports = main;
+
+// TODO(developer): Uncomment below lines before running the sample.
+// main(...process.argv.slice(2)).catch(err => {
+//   console.error(err);
+//   process.exitCode = 1;
+// });
diff --git a/tpu/queuedResources/getQueuedResource.js b/tpu/queuedResources/getQueuedResource.js
@@ -16,25 +16,27 @@
 
 'use strict';
 
-async function main(queuedResourceName, zone) {
+async function main(tpuClient) {
   // [START tpu_queued_resources_get]
-  // Import the TPU library
-  const {TpuClient} = require('@google-cloud/tpu').v2alpha1;
+  // Import the TPUClient
+  // TODO(developer): Uncomment below line before running the sample.
+  // const {TpuClient} = require('@google-cloud/tpu').v2alpha1;
 
   // Instantiate a tpuClient
-  const tpuClient = new TpuClient();
+  // TODO(developer): Uncomment below line before running the sample.
+  // tpuClient = new TpuClient();
 
   /**
-   * TODO(developer): Update/uncomment these variables before running the sample.
+   * TODO(developer): Update these variables before running the sample.
    */
   // Project ID or project number of the Google Cloud project, where you want to retrive node.
   const projectId = await tpuClient.getProjectId();
 
   // The name of queued resource.
-  // queuedResourceName = 'queued-resource-1';
+  const queuedResourceName = 'queued-resource-1';
 
   // The zone of your queued resource.
-  // zone = 'europe-west4-a';
+  const zone = 'europe-west4-a';
 
   async function callGetQueuedResource() {
     const request = {
@@ -43,14 +45,17 @@ async function main(queuedResourceName, zone) {
 
     const [response] = await tpuClient.getQueuedResource(request);
 
-    console.log(JSON.stringify(response));
     console.log(`Queued resource ${queuedResourceName} retrived.`);
+    return response;
   }
-  await callGetQueuedResource();
+  return await callGetQueuedResource();
   // [END tpu_queued_resources_get]
 }
 
-main(...process.argv.slice(2)).catch(err => {
-  console.error(err);
-  process.exitCode = 1;
-});
+module.exports = main;
+
+// TODO(developer): Uncomment below lines before running the sample.
+// main(...process.argv.slice(2)).catch(err => {
+//   console.error(err);
+//   process.exitCode = 1;
+// });
diff --git a/tpu/queuedResources/getQueuedResourcesList.js b/tpu/queuedResources/getQueuedResourcesList.js
@@ -16,13 +16,15 @@
 
 'use strict';
 
-async function main(zone) {
+async function main(tpuClient) {
   // [START tpu_queued_resources_list]
-  // Import the TPU library
-  const {TpuClient} = require('@google-cloud/tpu').v2alpha1;
+  // Import the TPUClient
+  // TODO(developer): Uncomment below line before running the sample.
+  // const {TpuClient} = require('@google-cloud/tpu').v2alpha1;
 
   // Instantiate a tpuClient
-  const tpuClient = new TpuClient();
+  // TODO(developer): Uncomment below line before running the sample.
+  // tpuClient = new TpuClient();
 
   /**
    * TODO(developer): Update/uncomment these variables before running the sample.
@@ -31,7 +33,7 @@ async function main(zone) {
   const projectId = await tpuClient.getProjectId();
 
   // The zone from which the Queued Resources are retrived.
-  // zone = 'europe-west4-a';
+  const zone = 'europe-west4-a';
 
   async function callGetQueuedResourcesList() {
     const request = {
@@ -40,13 +42,16 @@ async function main(zone) {
 
     const [response] = await tpuClient.listQueuedResources(request);
 
-    console.log(JSON.stringify(response));
+    return response;
   }
-  await callGetQueuedResourcesList();
+  return await callGetQueuedResourcesList();
   // [END tpu_queued_resources_list]
 }
 
-main(...process.argv.slice(2)).catch(err => {
-  console.error(err);
-  process.exitCode = 1;
-});
+module.exports = main;
+
+// TODO(developer): Uncomment below lines before running the sample.
+// main(...process.argv.slice(2)).catch(err => {
+//   console.error(err);
+//   process.exitCode = 1;
+// });