Use mocked TPUClient in tests

GoogleCloudPlatform · Dec 4, 2024 · 1a65640 · 1a65640
1 parent 0e86993
commit 1a65640
Show file tree

Hide file tree

Showing 8 changed files with 254 additions and 197 deletions.
diff --git a/tpu/package.json b/tpu/package.json
@@ -21,27 +21,4 @@
     "c8": "^10.0.0",
     "mocha": "^10.0.0"
   }
-}
-{
-  "name": "nodejs-docs-samples-tpu",
-  "license": "Apache-2.0",
-  "author": "Google Inc.",
-  "engines": {
-    "node": ">=16.0.0"
-  },
-  "repository": "googleapis/nodejs-tpu",
-  "private": true,
-  "files": [
-    "*.js"
-  ],
-  "scripts": {
-    "test": "c8 mocha -p -j 2 test --timeout 1200000"
-  },
-  "dependencies": {
-    "@google-cloud/tpu": "^3.5.0"
-  },
-  "devDependencies": {
-    "c8": "^10.0.0",
-    "mocha": "^10.0.0"
-  }
-}
+}
diff --git a/tpu/queuedResources/createQueuedResource.js b/tpu/queuedResources/createQueuedResource.js
@@ -16,53 +16,49 @@
 
 '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();
 
   // The name of the network you want the node to connect to. The network should be assigned to your project.
   const networkName = 'compute-tpu-network';
 
-  // The region of the network, that you want the node to connect to.
-  const region = 'europe-west4';
+  // The region of the subnetwork, that you want the node to connect to.
+  const region = 'us-central1';
 
   // 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 = `${region}-f`;
 
   // 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 = 'us-central1-f';
 
   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 = 'us-central1-f';
 
   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 = 'us-central1-f';
 
   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 = 'us-central1-f';
 
   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;
+// });