DLP: Updated test cases for storedInfoType related samples to use moc…

…k approach (#3414) Tests of following region tag samples updated -dlp_create_stored_infotype -dlp_update_stored_infotype
GoogleCloudPlatform · Aug 22, 2023 · 5383ced · 5383ced
1 parent 99e30d2
commit 5383ced
Show file tree

Hide file tree

Showing 4 changed files with 196 additions and 97 deletions.
diff --git a/dlp/createStoredInfoType.js b/dlp/createStoredInfoType.js
@@ -18,7 +18,7 @@
 // title: Create stored infotype.
 // description: Uses the Data Loss Prevention API to create a stored infotype.
 // usage: node createStoredInfoType.js projectId  infoTypeId, outputPath, dataProjectId, datasetId, tableId, fieldName
-function main(
+async function main(
   projectId,
   infoTypeId,
   outputPath,
@@ -99,7 +99,7 @@ function main(
     // Print results
     console.log(`InfoType stored successfully: ${response.name}`);
   }
-  createStoredInfoType();
+  await createStoredInfoType();
   // [END dlp_create_stored_infotype]
 }
 
@@ -108,4 +108,7 @@ process.on('unhandledRejection', err => {
   process.exitCode = 1;
 });
 
-main(...process.argv.slice(2));
+// TODO(developer): Please uncomment below line before running sample
+// main(...process.argv.slice(2));
+
+module.exports = main;
diff --git a/dlp/system-test/metadata.test.js b/dlp/system-test/metadata.test.js
@@ -20,11 +20,14 @@ const cp = require('child_process');
 const uuid = require('uuid');
 const DLP = require('@google-cloud/dlp');
 const {Storage} = require('@google-cloud/storage');
+const proxyquire = require('proxyquire');
+const sinon = require('sinon');
+
+const {MOCK_DATA} = require('./mockdata');
 
 const dataProject = 'bigquery-public-data';
 const dataSetId = 'samples';
 const tableId = 'github_nested';
-const fieldId = 'url';
 
 const storage = new Storage();
 const testFile = 'resources/test.txt';
@@ -35,7 +38,6 @@ const execSync = cmd => cp.execSync(cmd, {encoding: 'utf-8'});
 const client = new DLP.DlpServiceClient();
 describe('metadata', () => {
   let projectId, storedInfoTypeId;
-  const infoTypeCloudStorageFileSet = `gs://${bucketName}/test.txt`;
 
   before(async () => {
     projectId = await client.getProjectId();
@@ -58,6 +60,7 @@ describe('metadata', () => {
 
   // Delete stored infotypes created in the snippets.
   afterEach(async () => {
+    sinon.restore();
     if (!storedInfoTypeId) {
       return;
     }
@@ -85,108 +88,147 @@ describe('metadata', () => {
   });
 
   // dlp_create_stored_infotype
-  it('should create a stored infotype', () => {
-    const infoTypeId = `stored-infoType-${uuid.v4()}`;
-    const infoTypeOutputPath = `gs://${bucketName}`;
-    const output = execSync(
-      `node createStoredInfoType.js ${projectId} ${infoTypeId} ${infoTypeOutputPath} ${dataProject} ${dataSetId} ${tableId} ${fieldId}`
+  it('should create a stored infotype', async () => {
+    const storedInfoTypeName = 'MOCK_INFOTYPE';
+    const infoTypeId = 'MOCK_INFOTYPE_ID';
+    const outputPath = 'MOCK_OUTPUT_PATH';
+    const fieldName = 'MOCK_FIELD';
+    const DATA_CONSTANTS = MOCK_DATA.CREATE_STORED_INFOTYPE(
+      projectId,
+      infoTypeId,
+      outputPath,
+      dataProject,
+      dataSetId,
+      tableId,
+      fieldName,
+      storedInfoTypeName
+    );
+    const mockCreateStoredInfoType = sinon
+      .stub()
+      .resolves([{name: storedInfoTypeName}]);
+    sinon.replace(
+      DLP.DlpServiceClient.prototype,
+      'createStoredInfoType',
+      mockCreateStoredInfoType
+    );
+
+    const mockConsoleLog = sinon.stub();
+    sinon.replace(console, 'log', mockConsoleLog);
+
+    const createStoredInfoType = proxyquire('../createStoredInfoType', {
+      '@google-cloud/dlp': {DLP: DLP},
+    });
+    await createStoredInfoType(
+      projectId,
+      infoTypeId,
+      outputPath,
+      dataProject,
+      dataSetId,
+      tableId,
+      fieldName
+    );
+    sinon.assert.calledOnceWithExactly(
+      mockCreateStoredInfoType,
+      DATA_CONSTANTS.REQUEST_CREATE_DLP_JOB
+    );
+    sinon.assert.calledWithExactly(
+      mockConsoleLog,
+      `InfoType stored successfully: ${storedInfoTypeName}`
     );
-    assert.match(output, /InfoType stored successfully:/);
-    storedInfoTypeId = output.split(':')[1].trim();
   });
 
-  it('should handle stored infotype creation errors', () => {
-    let output;
-    const infoTypeId = `stored-infoType-${uuid.v4()}`;
-    const infoTypeOutputPath = 'INFOTYPE_OUTPUT_PATH';
+  it('should handle error while creating stored infotype', async () => {
+    const infoTypeId = 'MOCK_INFOTYPE_ID';
+    const outputPath = 'MOCK_OUTPUT_PATH';
+    const fieldName = 'MOCK_FIELD';
+    const mockCreateStoredInfoType = sinon.stub().rejects(new Error('Failed'));
+    sinon.replace(
+      DLP.DlpServiceClient.prototype,
+      'createStoredInfoType',
+      mockCreateStoredInfoType
+    );
+
+    const mockConsoleLog = sinon.stub();
+    sinon.replace(console, 'log', mockConsoleLog);
+
+    const createStoredInfoType = proxyquire('../createStoredInfoType', {
+      '@google-cloud/dlp': {DLP: DLP},
+    });
     try {
-      output = execSync(
-        `node createStoredInfoType.js BAD_PROJECT_ID ${infoTypeId} ${infoTypeOutputPath} ${dataProject} ${dataSetId} ${tableId} ${fieldId}`
+      await createStoredInfoType(
+        projectId,
+        infoTypeId,
+        outputPath,
+        dataProject,
+        dataSetId,
+        tableId,
+        fieldName
       );
-    } catch (err) {
-      output = err.message;
+    } catch (error) {
+      assert.equal(error.message, 'Failed');
     }
-    assert.include(output, 'INVALID_ARGUMENT');
   });
 
   // dlp_update_stored_infotype
   it('should update a stored infotype', async () => {
-    let output;
-    const infoTypeId = `stored-infoType-${uuid.v4()}`;
-    const infoTypeOutputPath = `gs://${bucketName}`;
-    try {
-      // First create a temporary stored infoType
-      const [response] = await client.createStoredInfoType({
-        parent: `projects/${projectId}/locations/global`,
-        config: {
-          displayName: 'GitHub usernames',
-          description: 'Dictionary of GitHub usernames used in commits',
-          largeCustomDictionary: {
-            outputPath: {
-              path: infoTypeOutputPath,
-            },
-            bigQueryField: {
-              table: {
-                datasetId: dataSetId,
-                projectId: dataProject,
-                tableId: tableId,
-              },
-              field: {
-                name: fieldId,
-              },
-            },
-          },
-        },
-        storedInfoTypeId: infoTypeId,
-      });
-      storedInfoTypeId = response.name;
-      // Execute the update script
-      output = execSync(
-        `node updateStoredInfoType.js ${projectId} ${infoTypeId} ${infoTypeOutputPath} ${infoTypeCloudStorageFileSet}`
-      );
-    } catch (err) {
-      output = err.message;
-    }
-    assert.match(output, /InfoType updated successfully:/);
+    const storedInfoTypeName = 'MOCK_INFOTYPE';
+    const infoTypeId = 'MOCK_INFOTYPE_ID';
+    const outputPath = 'MOCK_OUTPUT_PATH';
+    const fileSetUrl = 'MOCK_FILE_SET_URL';
+    const DATA_CONSTANTS = MOCK_DATA.UPDATE_STORED_INFOTYPE(
+      projectId,
+      infoTypeId,
+      outputPath,
+      fileSetUrl
+    );
+    const mockUpdateStoredInfoType = sinon
+      .stub()
+      .resolves([{name: storedInfoTypeName}]);
+    sinon.replace(
+      DLP.DlpServiceClient.prototype,
+      'updateStoredInfoType',
+      mockUpdateStoredInfoType
+    );
+
+    const mockConsoleLog = sinon.stub();
+    sinon.replace(console, 'log', mockConsoleLog);
+
+    const updateStoredInfoType = proxyquire('../updateStoredInfoType', {
+      '@google-cloud/dlp': {DLP: DLP},
+    });
+    await updateStoredInfoType(projectId, infoTypeId, outputPath, fileSetUrl);
+
+    sinon.assert.calledWith(
+      mockUpdateStoredInfoType,
+      DATA_CONSTANTS.REQUEST_UPDATE_STORED_INFOTYPE
+    );
+    sinon.assert.calledWithMatch(
+      mockConsoleLog,
+      'InfoType updated successfully:'
+    );
   });
 
-  it('should handle stored infotype update errors', async () => {
-    let output;
-    const infoTypeId = `stored-infoType-${uuid.v4()}`;
-    const infoTypeOutputPath = 'INFOTYPE_OUTPUT_PATH';
+  it('should handle error while updating stored infotype', async () => {
+    const infoTypeId = 'MOCK_INFOTYPE_ID';
+    const outputPath = 'MOCK_OUTPUT_PATH';
+    const fileSetUrl = 'MOCK_FILE_SET_URL';
+    const mockUpdateStoredInfoType = sinon.stub().rejects(new Error('Failed'));
+    sinon.replace(
+      DLP.DlpServiceClient.prototype,
+      'updateStoredInfoType',
+      mockUpdateStoredInfoType
+    );
+
+    const mockConsoleLog = sinon.stub();
+    sinon.replace(console, 'log', mockConsoleLog);
+
+    const updateStoredInfoType = proxyquire('../updateStoredInfoType', {
+      '@google-cloud/dlp': {DLP: DLP},
+    });
     try {
-      // First create a temporary stored infoType
-      const [response] = await client.createStoredInfoType({
-        parent: `projects/${projectId}/locations/global`,
-        config: {
-          displayName: 'GitHub usernames',
-          description: 'Dictionary of GitHub usernames used in commits',
-          largeCustomDictionary: {
-            outputPath: {
-              path: infoTypeOutputPath,
-            },
-            bigQueryField: {
-              table: {
-                datasetId: dataSetId,
-                projectId: dataProject,
-                tableId: tableId,
-              },
-              field: {
-                name: fieldId,
-              },
-            },
-          },
-        },
-        storedInfoTypeId: infoTypeId,
-      });
-      storedInfoTypeId = response.name;
-      // Execute the update script
-      output = execSync(
-        `node updateStoredInfoType.js BAD_PROJECT_ID ${infoTypeId} ${infoTypeOutputPath} ${infoTypeCloudStorageFileSet}`
-      );
-    } catch (err) {
-      output = err.message;
+      await updateStoredInfoType(projectId, infoTypeId, outputPath, fileSetUrl);
+    } catch (error) {
+      assert.equal(error.message, 'Failed');
     }
-    assert.include(output, 'INVALID_ARGUMENT');
   });
 });
diff --git a/dlp/system-test/mockdata.js b/dlp/system-test/mockdata.js
@@ -1197,6 +1197,57 @@ const MOCK_DATA = {
       nack: sinon.stub(),
     },
   }),
+  CREATE_STORED_INFOTYPE: (
+    projectId,
+    infoTypeId,
+    outputPath,
+    dataProjectId,
+    datasetId,
+    tableId,
+    fieldName
+  ) => ({
+    REQUEST_CREATE_DLP_JOB: {
+      parent: `projects/${projectId}/locations/global`,
+      config: {
+        displayName: 'GitHub usernames',
+        description: 'Dictionary of GitHub usernames used in commits',
+        largeCustomDictionary: {
+          outputPath: {
+            path: outputPath,
+          },
+          bigQueryField: {
+            table: {
+              datasetId: datasetId,
+              projectId: dataProjectId,
+              tableId: tableId,
+            },
+            field: {
+              name: fieldName,
+            },
+          },
+        },
+      },
+      storedInfoTypeId: infoTypeId,
+    },
+  }),
+  UPDATE_STORED_INFOTYPE: (projectId, infoTypeId, outputPath, fileSetUrl) => ({
+    REQUEST_UPDATE_STORED_INFOTYPE: {
+      name: `projects/${projectId}/storedInfoTypes/${infoTypeId}`,
+      config: {
+        largeCustomDictionary: {
+          outputPath: {
+            path: outputPath,
+          },
+          cloudStorageFileSet: {
+            url: fileSetUrl,
+          },
+        },
+      },
+      updateMask: {
+        paths: ['large_custom_dictionary.cloud_storage_file_set.url'],
+      },
+    },
+  }),
 };
 
 module.exports = {MOCK_DATA};
diff --git a/dlp/updateStoredInfoType.js b/dlp/updateStoredInfoType.js
@@ -18,7 +18,7 @@
 //  title: Update stored infoType.
 //  description: Uses the Data Loss Prevention API to update a stored infoType.
 //  usage: node updateStoredInfoType.js projectId  infoTypeId, outputPath, fileSetUrl
-function main(projectId, infoTypeId, outputPath, fileSetUrl) {
+async function main(projectId, infoTypeId, outputPath, fileSetUrl) {
   // [START dlp_update_stored_infotype]
   // Import the required libraries
   const dlp = require('@google-cloud/dlp');
@@ -68,7 +68,7 @@ function main(projectId, infoTypeId, outputPath, fileSetUrl) {
     // Print the results.
     console.log(`InfoType updated successfully: ${JSON.stringify(response)}`);
   }
-  updateStoredInfoType();
+  await updateStoredInfoType();
   // [END dlp_update_stored_infotype]
 }
 
@@ -77,4 +77,7 @@ process.on('unhandledRejection', err => {
   process.exitCode = 1;
 });
 
-main(...process.argv.slice(2));
+// TODO(developer): Please uncomment below line before running sample
+// main(...process.argv.slice(2));
+
+module.exports = main;