Use mocked TPUClient in tests

GoogleCloudPlatform · Dec 2, 2024 · 1a62542 · 1a62542
1 parent debf398
commit 1a62542
Show file tree

Hide file tree

Showing 8 changed files with 213 additions and 239 deletions.
diff --git a/tpu/package.json b/tpu/package.json
@@ -3,21 +3,22 @@
     "license": "Apache-2.0",
     "author": "Google Inc.",
     "engines": {
-      "node": ">=16.0.0"
+        "node": ">=16.0.0"
     },
     "repository": "googleapis/nodejs-tpu",
     "private": true,
     "files": [
-      "*.js"
+        "*.js"
     ],
     "scripts": {
-      "test": "c8 mocha -p -j 2 test --timeout 1200000"
+        "test": "c8 mocha -p -j 2 test --timeout 1200000"
     },
     "dependencies": {
-      "@google-cloud/tpu": "^3.5.0"
+        "@google-cloud/tpu": "^3.5.0",
+        "sinon": "^19.0.2"
     },
     "devDependencies": {
-      "c8": "^10.0.0",
-      "mocha": "^10.0.0"
+        "c8": "^10.0.0",
+        "mocha": "^10.0.0"
     }
-  }
+}
diff --git a/tpu/queuedResources/createQueuedResourceNetwork.js b/tpu/queuedResources/createQueuedResourceNetwork.js
@@ -16,21 +16,17 @@
 
 'use strict';
 
-async function main(
-  nodeName,
-  queuedResourceName,
-  zone,
-  tpuType,
-  tpuSoftwareVersion
-) {
-  // [START tpu_queued_resources_create_network]
-  // Import the TPU library
-  const {TpuClient} = require('@google-cloud/tpu').v2alpha1;
+async function main(tpuClient) {
+  // [START tpu_queued_resources_network]
+  // 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.
@@ -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 callCreateQueuedResourceNetwork() {
     // Specify the network and subnetwork that you want to connect your TPU to.
@@ -112,13 +108,16 @@ async function main(
 
     // You can wait until TPU Node is READY,
     // and check its status using getTpuVm() from `tpu_vm_get` sample.
-    console.log(JSON.stringify(response));
+    return response;
   }
-  await callCreateQueuedResourceNetwork();
-  // [END tpu_queued_resources_create_network]
+  return await callCreateQueuedResourceNetwork();
+  // [END tpu_queued_resources_network]
 }
 
-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/createQueuedResourceStartupScript.js b/tpu/queuedResources/createQueuedResourceStartupScript.js
@@ -16,21 +16,17 @@
 
 'use strict';
 
-async function main(
-  nodeName,
-  queuedResourceName,
-  zone,
-  tpuType,
-  tpuSoftwareVersion
-) {
+async function main(tpuClient) {
   // [START tpu_queued_resources_startup_script]
-  // 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.
@@ -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 callCreateQueuedResourceStartupScript() {
     // Create a node
@@ -116,13 +112,16 @@ async function main(
 
     // You can wait until TPU Node is READY,
     // and check its status using getTpuVm() from `tpu_vm_get` sample.
-    console.log(JSON.stringify(response));
+    return response;
   }
-  await callCreateQueuedResourceStartupScript();
+  return await callCreateQueuedResourceStartupScript();
   // [END tpu_queued_resources_startup_script]
 }
 
-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/createQueuedResourceTimeBound.js b/tpu/queuedResources/createQueuedResourceTimeBound.js
@@ -16,21 +16,17 @@
 
 'use strict';
 
-async function main(
-  nodeName,
-  queuedResourceName,
-  zone,
-  tpuType,
-  tpuSoftwareVersion
-) {
+async function main(tpuClient) {
   // [START tpu_queued_resources_time_bound]
-  // 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.
@@ -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 callCreateQueuedResourceTimeBound() {
     // Create a node
@@ -143,13 +139,16 @@ async function main(
 
     // You can wait until TPU Node is READY,
     // and check its status using getTpuVm() from `tpu_vm_get` sample.
-    console.log(JSON.stringify(response));
+    return response;
   }
-  await callCreateQueuedResourceTimeBound();
+  return await callCreateQueuedResourceTimeBound();
   // [END tpu_queued_resources_time_bound]
 }
 
-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