From a9ee1f7ab5cd8f1417285dcff8a97ed615d01c47 Mon Sep 17 00:00:00 2001
From: Joanna Grycz <jgrycz@softserveinc.com>
Date: Tue, 12 Nov 2024 18:05:46 +0100
Subject: [PATCH] Use mocked TPUClient in tests
---
.../createQueuedResourceNetwork.js | 47 +++++-----
.../createQueuedResourceStartupScript.js | 43 +++++----
.../createQueuedResourceTimeBound.js | 43 +++++----
.../forceDeleteQueuedResource.js | 59 ------------
tpu/test/createQueuedResourceNetwork.test.js | 91 ++++++++++---------
.../createQueuedResourceStartupScript.test.js | 84 ++++++++++-------
.../createQueuedResourceTimeBound.test.js | 76 +++++++++-------
7 files changed, 208 insertions(+), 235 deletions(-)
delete mode 100644 tpu/queuedResources/forceDeleteQueuedResource.js
diff --git a/tpu/queuedResources/createQueuedResourceNetwork.js b/tpu/queuedResources/createQueuedResourceNetwork.js
index 1ec9a9b7ed..2b36291d3a 100644
--- 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.
@@ -42,27 +38,27 @@ async function main(
const networkName = 'compute-tpu-network';
// The region of the network, that you want the node to connect to.
- const region = 'europe-west4';
+ 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 = `${zone}-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
index ce0c2080ca..7e5b499c1f 100644
--- 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.
@@ -42,27 +38,27 @@ async function main(
const networkName = 'compute-tpu-network';
// The region of the network, that you want the node to connect to.
- const region = 'europe-west4';
+ 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 = `${zone}-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
index 92e1889c2d..535d8d5803 100644
--- 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.
@@ -42,27 +38,27 @@ async function main(
const networkName = 'compute-tpu-network';
// The region of the network, that you want the node to connect to.
- const region = 'europe-west4';
+ 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 = `${zone}-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
deleted file mode 100644
index e842c39e26..0000000000
--- a/tpu/queuedResources/forceDeleteQueuedResource.js
+++ /dev/null
@@ -1,59 +0,0 @@
-/*
- * Copyright 2024 Google LLC
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- * https://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-'use strict';
-
-async function main(queuedResourceName, zone) {
- // [START tpu_queued_resources_delete_force]
- // Import the TPU library
- const {TpuClient} = require('@google-cloud/tpu').v2alpha1;
-
- // Instantiate a tpuClient
- const tpuClient = new TpuClient();
-
- /**
- * TODO(developer): Update/uncomment 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';
-
- // The zone of your queued resource.
- // zone = 'europe-west4-a';
-
- async function callForceDeleteQueuedResource() {
- const request = {
- name: `projects/${projectId}/locations/${zone}/queuedResources/${queuedResourceName}`,
- force: true,
- };
-
- const [operation] = await tpuClient.deleteQueuedResource(request);
-
- // Wait for the delete operation to complete.
- await operation.promise();
-
- console.log(`Queued resource ${queuedResourceName} deletion forced.`);
- }
- await callForceDeleteQueuedResource();
- // [END tpu_queued_resources_delete_force]
-}
-
-main(...process.argv.slice(2)).catch(err => {
- console.error(err);
- process.exitCode = 1;
-});
diff --git a/tpu/test/createQueuedResourceNetwork.test.js b/tpu/test/createQueuedResourceNetwork.test.js
index 898ffe9be1..3c18558794 100644
--- a/tpu/test/createQueuedResourceNetwork.test.js
+++ b/tpu/test/createQueuedResourceNetwork.test.js
@@ -16,57 +16,66 @@
'use strict';
-const path = require('path');
const assert = require('node:assert/strict');
-const {after, before, describe, it} = require('mocha');
-const cp = require('child_process');
-const {TpuClient} = require('@google-cloud/tpu').v2alpha1;
-
-const execSync = cmd => cp.execSync(cmd, {encoding: 'utf-8'});
-const cwd = path.join(__dirname, '..');
+const {beforeEach, afterEach, describe, it} = require('mocha');
+const sinon = require('sinon');
+const createQueuedResourceNetwork = require('../queuedResources/createQueuedResourceNetwork.js');
describe('TPU queued resource with specified network', async () => {
- const queuedResourceName = `queued-resource-with-network-${Math.floor(Math.random() * 1000 + 1)}`;
- const nodeName = `node-with-network-2a2b3c${Math.floor(Math.random() * 1000 + 1)}`;
- const zone = 'us-south1-a';
- const tpuType = 'v5litepod-1';
- const tpuSoftwareVersion = 'tpu-vm-tf-2.14.1';
- let projectId;
+ const queuedResourceName = 'queued-resource-1';
+ const nodeName = 'node-name-1';
+ const zone = 'us-central1-a';
+ const projectId = 'project_id';
+ let tpuClientMock;
- before(async () => {
- const tpuClient = new TpuClient();
- projectId = await tpuClient.getProjectId();
+ beforeEach(() => {
+ tpuClientMock = {
+ getProjectId: sinon.stub().resolves(projectId),
+ };
});
- after(() => {
- // Delete queued resource
- execSync(
- `node ./queuedResources/forceDeleteQueuedResource.js ${queuedResourceName} ${zone}`,
- {
- cwd,
- }
- );
+ afterEach(() => {
+ sinon.restore();
});
- it('should create queued resource', () => {
- const networkConfig = {
- network: `projects/${projectId}/global/networks/compute-tpu-network`,
- subnetwork: `projects/${projectId}/regions/europe-west4/subnetworks/compute-tpu-network`,
- enableExternalIps: true,
- };
+ it('should create queued resource', async () => {
+ tpuClientMock.createQueuedResource = sinon.stub().resolves([
+ {
+ promise: sinon.stub().resolves([
+ {
+ name: queuedResourceName,
+ },
+ ]),
+ },
+ ]);
- const response = JSON.parse(
- execSync(
- `node ./queuedResources/createQueuedResourceNetwork.js ${nodeName} ${queuedResourceName} ${zone} ${tpuType} ${tpuSoftwareVersion}`,
- {
- cwd,
- }
- )
- );
+ const response = await createQueuedResourceNetwork(tpuClientMock);
- assert.deepEqual(
- response.tpu.nodeSpec[0].node.networkConfig,
- networkConfig
+ sinon.assert.calledWith(
+ tpuClientMock.createQueuedResource,
+ sinon.match({
+ parent: `projects/${projectId}/locations/${zone}`,
+ queuedResource: {
+ name: queuedResourceName,
+ tpu: {
+ nodeSpec: [
+ sinon.match({
+ parent: `projects/${projectId}/locations/${zone}`,
+ node: {
+ networkConfig: {
+ network: `projects/${projectId}/global/networks/compute-tpu-network`,
+ subnetwork: `projects/${projectId}/regions/europe-west4/subnetworks/compute-tpu-network`,
+ enableExternalIps: true,
+ },
+ },
+ nodeId: nodeName,
+ }),
+ ],
+ },
+ },
+ queuedResourceId: queuedResourceName,
+ })
);
+ assert(response.name.includes(queuedResourceName));
});
});
diff --git a/tpu/test/createQueuedResourceStartupScript.test.js b/tpu/test/createQueuedResourceStartupScript.test.js
index abf368f467..60c587e303 100644
--- a/tpu/test/createQueuedResourceStartupScript.test.js
+++ b/tpu/test/createQueuedResourceStartupScript.test.js
@@ -16,46 +16,62 @@
'use strict';
-const path = require('path');
const assert = require('node:assert/strict');
-const {after, describe, it} = require('mocha');
-const cp = require('child_process');
-
-const execSync = cmd => cp.execSync(cmd, {encoding: 'utf-8'});
-const cwd = path.join(__dirname, '..');
+const {beforeEach, afterEach, describe, it} = require('mocha');
+const sinon = require('sinon');
+const createQueuedResourceStartupScript = require('../queuedResources/createQueuedResourceStartupScript.js');
describe('TPU queued resource with start-up script', async () => {
- const queuedResourceName = `queued-resource-startup-script-${Math.floor(Math.random() * 1000 + 1)}`;
- const nodeName = `node-startup-script-2a2b3c${Math.floor(Math.random() * 1000 + 1)}`;
- const zone = 'us-east1-d';
- const tpuType = 'v3-32';
- const tpuSoftwareVersion = 'tpu-vm-tf-2.14.1';
-
- after(() => {
- // Delete queued resource
- execSync(
- `node ./queuedResources/forceDeleteQueuedResource.js ${queuedResourceName} ${zone}`,
- {
- cwd,
- }
- );
- });
+ const queuedResourceName = 'queued-resource-1';
+ const nodeName = 'node-name-1';
+ const zone = 'us-central1-a';
+ const projectId = 'project_id';
+ let tpuClientMock;
- it('should create queued resource', () => {
- const metadata = {
- 'startup-script':
- '#!/bin/bash\n echo "Hello World" > /var/log/hello.log\n sudo pip3 install --upgrade numpy >> /var/log/hello.log 2>&1',
+ beforeEach(() => {
+ tpuClientMock = {
+ getProjectId: sinon.stub().resolves(projectId),
};
+ });
+
+ afterEach(() => {
+ sinon.restore();
+ });
+
+ it('should create queued resource', async () => {
+ tpuClientMock.createQueuedResource = sinon.stub().resolves([
+ {
+ promise: sinon.stub().resolves([
+ {
+ name: queuedResourceName,
+ },
+ ]),
+ },
+ ]);
- const response = JSON.parse(
- execSync(
- `node ./queuedResources/createQueuedResourceStartupScript.js ${nodeName} ${queuedResourceName} ${zone} ${tpuType} ${tpuSoftwareVersion}`,
- {
- cwd,
- }
- )
- );
+ const response = await createQueuedResourceStartupScript(tpuClientMock);
- assert.deepEqual(response.tpu.nodeSpec[0].node.metadata, metadata);
+ sinon.match({
+ parent: `projects/${projectId}/locations/${zone}`,
+ queuedResource: {
+ name: queuedResourceName,
+ tpu: {
+ nodeSpec: [
+ {
+ parent: `projects/${projectId}/locations/${zone}`,
+ node: {
+ metadata: {
+ 'startup-script':
+ '#!/bin/bash\n echo "Hello World" > /var/log/hello.log\n sudo pip3 install --upgrade numpy >> /var/log/hello.log 2>&1',
+ },
+ },
+ nodeId: nodeName,
+ },
+ ],
+ },
+ },
+ queuedResourceId: queuedResourceName,
+ });
+ assert(response.name.includes(queuedResourceName));
});
});
diff --git a/tpu/test/createQueuedResourceTimeBound.test.js b/tpu/test/createQueuedResourceTimeBound.test.js
index 2a70a2dcd3..43b6dc8d2f 100644
--- a/tpu/test/createQueuedResourceTimeBound.test.js
+++ b/tpu/test/createQueuedResourceTimeBound.test.js
@@ -16,44 +16,54 @@
'use strict';
-const path = require('path');
const assert = require('node:assert/strict');
-const {after, describe, it} = require('mocha');
-const cp = require('child_process');
-
-const execSync = cmd => cp.execSync(cmd, {encoding: 'utf-8'});
-const cwd = path.join(__dirname, '..');
+const {beforeEach, afterEach, describe, it} = require('mocha');
+const sinon = require('sinon');
+const createQueuedResourceTimeBound = require('../queuedResources/createQueuedResourceTimeBound.js');
describe('TPU time bound queued resource', async () => {
- const queuedResourceName = `queued-resource-time-bound-${Math.floor(Math.random() * 1000 + 1)}`;
- const nodeName = `node-time-bound-2a2b3c${Math.floor(Math.random() * 1000 + 1)}`;
- const zone = 'us-west4-a';
- const tpuType = 'v5litepod-1';
- const tpuSoftwareVersion = 'tpu-vm-tf-2.14.1';
-
- after(() => {
- // Delete queued resource
- execSync(
- `node ./queuedResources/forceDeleteQueuedResource.js ${queuedResourceName} ${zone}`,
- {
- cwd,
- }
- );
+ const queuedResourceName = 'queued-resource-1';
+ const zone = 'us-central1-a';
+ const projectId = 'project_id';
+ let tpuClientMock;
+
+ beforeEach(() => {
+ tpuClientMock = {
+ getProjectId: sinon.stub().resolves(projectId),
+ };
});
- it('should create queued resource', () => {
- const response = JSON.parse(
- execSync(
- `node ./queuedResources/createQueuedResourceTimeBound.js ${nodeName} ${queuedResourceName} ${zone} ${tpuType} ${tpuSoftwareVersion}`,
- {
- cwd,
- }
- )
- );
+ afterEach(() => {
+ sinon.restore();
+ });
+
+ it('should create queued resource', async () => {
+ tpuClientMock.createQueuedResource = sinon.stub().resolves([
+ {
+ promise: sinon.stub().resolves([
+ {
+ name: queuedResourceName,
+ },
+ ]),
+ },
+ ]);
- assert.ok(response.queueingPolicy);
- assert.ok(response.queueingPolicy.validAfterTime);
- assert(typeof response.queueingPolicy.validAfterTime.seconds, 'string');
- assert(typeof response.queueingPolicy.validAfterTime.nano, 'number');
+ const response = await createQueuedResourceTimeBound(tpuClientMock);
+
+ sinon.assert.calledWith(
+ tpuClientMock.createQueuedResource,
+ sinon.match({
+ parent: `projects/${projectId}/locations/${zone}`,
+ queuedResource: {
+ queueingPolicy: {
+ validAfterDuration: {
+ seconds: 6 * 3600,
+ },
+ },
+ },
+ queuedResourceId: 'queued-resource-1',
+ })
+ );
+ assert(response.name.includes(queuedResourceName));
});
});