chore(modelarmor): fixed floorsetting clean up and refactoring of tests

Author: mihirvala08

model-armor/snippets/getFolderFloorSettings.js

@@ -45,3 +45,15 @@ async function getFolderFloorSettings(folderId) {
 }
 
 module.exports = getFolderFloorSettings;
+
+/* c8 ignore next 10 */
+if (require.main === module) {
+  main(...process.argv.slice(2)).catch(err => {
+    console.error(err.message);
+    process.exitCode = 1;
+  });
+  process.on('unhandledRejection', err => {
+    console.error(err.message);
+    process.exitCode = 1;
+  });
+} 

model-armor/snippets/getOrganizationFloorSettings.js

@@ -47,3 +47,15 @@ async function getOrganizationFloorSettings(organizationId) {
 }
 
 module.exports = getOrganizationFloorSettings;
+
+/* c8 ignore next 10 */
+if (require.main === module) {
+  main(...process.argv.slice(2)).catch(err => {
+    console.error(err.message);
+    process.exitCode = 1;
+  });
+  process.on('unhandledRejection', err => {
+    console.error(err.message);
+    process.exitCode = 1;
+  });
+} 

model-armor/snippets/getProjectFloorSettings.js

@@ -47,3 +47,15 @@ async function getProjectFloorSettings(projectId) {
 }
 
 module.exports = getProjectFloorSettings;
+
+/* c8 ignore next 10 */
+if (require.main === module) {
+  main(...process.argv.slice(2)).catch(err => {
+    console.error(err.message);
+    process.exitCode = 1;
+  });
+  process.on('unhandledRejection', err => {
+    console.error(err.message);
+    process.exitCode = 1;
+  });
+} 

model-armor/snippets/updateFolderFloorSettings.js

@@ -73,3 +73,15 @@ async function updateFolderFloorSettings(folderId) {
 }
 
 module.exports = updateFolderFloorSettings;
+
+/* c8 ignore next 10 */
+if (require.main === module) {
+  main(...process.argv.slice(2)).catch(err => {
+    console.error(err.message);
+    process.exitCode = 1;
+  });
+  process.on('unhandledRejection', err => {
+    console.error(err.message);
+    process.exitCode = 1;
+  });
+} 

model-armor/snippets/updateOrganizationFloorSettings.js

@@ -71,3 +71,15 @@ async function updateOrganizationFloorSettings(organizationId) {
 }
 
 module.exports = updateOrganizationFloorSettings;
+
+/* c8 ignore next 10 */
+if (require.main === module) {
+  main(...process.argv.slice(2)).catch(err => {
+    console.error(err.message);
+    process.exitCode = 1;
+  });
+  process.on('unhandledRejection', err => {
+    console.error(err.message);
+    process.exitCode = 1;
+  });
+} 

model-armor/snippets/updateProjectFloorSettings.js

@@ -71,3 +71,15 @@ async function updateProjectFloorSettings(projectId) {
 }
 
 module.exports = updateProjectFloorSettings;
+
+/* c8 ignore next 10 */
+if (require.main === module) {
+  main(...process.argv.slice(2)).catch(err => {
+    console.error(err.message);
+    process.exitCode = 1;
+  });
+  process.on('unhandledRejection', err => {
+    console.error(err.message);
+    process.exitCode = 1;
+  });
+} 

model-armor/test/modelarmor.test.js

@@ -14,11 +14,10 @@
 
 'use strict';
 
-const { assert } = require('chai');
-const cp = require('child_process');
-const { v4: uuidv4 } = require('uuid');
-const { ModelArmorClient } = require('@google-cloud/modelarmor').v1;
-const { DlpServiceClient } = require('@google-cloud/dlp');
+const {assert} = require('chai');
+const {v4: uuidv4} = require('uuid');
+const {ModelArmorClient} = require('@google-cloud/modelarmor').v1;
+const {DlpServiceClient} = require('@google-cloud/dlp');
 
 let projectId;
 const locationId = process.env.GCLOUD_LOCATION || 'us-central1';
@@ -30,7 +29,6 @@ const options = {
 
 const client = new ModelArmorClient(options);
 const templateIdPrefix = `test-template-${uuidv4().substring(0, 8)}`;
-const execSync = cmd => cp.execSync(cmd, { encoding: 'utf-8' });
 
 let emptyTemplateId;
 let basicTemplateId;
@@ -79,15 +77,19 @@ async function deleteTemplate(templateName) {
   }
 }
 
+// eslint-disable-next-line no-unused-vars
 async function disableFloorSettings() {
   try {
+    // Create global client
+    const global_client = new ModelArmorClient();
+
     // Disable project floor settings
-    const [projectFloorSettings] = await client.getFloorSetting({
+    const [projectFloorSettings] = await global_client.getFloorSetting({
       name: `projects/${projectId}/locations/global/floorSetting`,
     });
 
     if (projectFloorSettings.enableFloorSettingEnforcement) {
-      const [updatedProjectSettings] = await client.updateFloorSetting({
+      const [updatedProjectSettings] = await global_client.updateFloorSetting({
         floorSetting: {
           name: `projects/${projectId}/locations/global/floorSetting`,
           enableFloorSettingEnforcement: false,
@@ -104,12 +106,12 @@ async function disableFloorSettings() {
 
     // Disable folder floor settings if folderId is available
     if (folderId) {
-      const [folderFloorSettings] = await client.getFloorSetting({
+      const [folderFloorSettings] = await global_client.getFloorSetting({
         name: `folders/${folderId}/locations/global/floorSetting`,
       });
 
       if (folderFloorSettings.enableFloorSettingEnforcement) {
-        const [updatedFolderSettings] = await client.updateFloorSetting({
+        const [updatedFolderSettings] = await global_client.updateFloorSetting({
           floorSetting: {
             name: `folders/${folderId}/locations/global/floorSetting`,
             enableFloorSettingEnforcement: false,
@@ -127,12 +129,12 @@ async function disableFloorSettings() {
 
     // Disable organization floor settings if organizationId is available
     if (organizationId) {
-      const [orgFloorSettings] = await client.getFloorSetting({
+      const [orgFloorSettings] = await global_client.getFloorSetting({
         name: `organizations/${organizationId}/locations/global/floorSetting`,
       });
 
       if (orgFloorSettings.enableFloorSettingEnforcement) {
-        const [updatedOrgSettings] = await client.updateFloorSetting({
+        const [updatedOrgSettings] = await global_client.updateFloorSetting({
           floorSetting: {
             name: `organizations/${organizationId}/locations/global/floorSetting`,
             enableFloorSettingEnforcement: false,
@@ -369,12 +371,10 @@ describe('Model Armor tests', () => {
 
   after(async () => {
     for (const templateName of templatesToDelete) {
-
-      // TODO: Enable clean up once tests are enabled.
+      // TODO(b/424365799): Uncomment below code once the mentioned issue is resolved
       // Disable floor settings to restore original state
       // await disableFloorSettings();
 
-
       await deleteTemplate(templateName);
     }
 
@@ -383,60 +383,58 @@ describe('Model Armor tests', () => {
 
   // =================== Floor Settings Tests ===================
 
-  // TODO: Enable these tests once floor settings API issue is resolved.
+  // TODO(b/424365799): Enable below tests once the mentioned issue is resolved
 
   it.skip('should get organization floor settings', () => {
-    const output = execSync(
-      `node snippets/getOrganizationFloorSettings.js ${organizationId}`
-    ).toString();
+    const getOrganizationFloorSettings = require('../snippets/getOrganizationFloorSettings');
+
+    const output = getOrganizationFloorSettings(organizationId).toString();
+
     const expectedName = `organizations/${organizationId}/locations/global/floorSetting`;
     assert.match(output, new RegExp(expectedName.replace(/\//g, '\\/')));
   });
 
   it.skip('should get folder floor settings', () => {
-    const output = execSync(
-      `node snippets/getFolderFloorSettings.js ${folderId}`
-    ).toString();
+    const getFolderFloorSettings = require('../snippets/getFolderFloorSettings');
+
+    const output = getFolderFloorSettings(folderId).toString();
 
     // Check for expected name format in output
     const expectedName = `folders/${folderId}/locations/global/floorSetting`;
     assert.match(output, new RegExp(expectedName.replace(/\//g, '\\/')));
   });
 
   it.skip('should get project floor settings', () => {
-    const output = execSync(
-      `node snippets/getProjectFloorSettings.js ${projectId}`
-    ).toString();
+    const getProjectFloorSettings = require('../snippets/getProjectFloorSettings');
+
+    const output = getProjectFloorSettings(projectId).toString();
 
     // Check for expected name format in output
     const expectedName = `projects/${projectId}/locations/global/floorSetting`;
     assert.match(output, new RegExp(expectedName.replace(/\//g, '\\/')));
   });
 
   it.skip('should update organization floor settings', () => {
-    const output = execSync(
-      `node snippets/updateOrganizationFloorSettings.js ${organizationId}`
-    ).toString();
+    const updateOrganizationFloorSettings = require('../snippets/updateOrganizationFloorSettings');
+    const output = updateOrganizationFloorSettings(organizationId).toString();
     // Check that the update was performed
     assert.match(output, /Updated organization floor settings/);
     // Check that the response contains enableFloorSettingEnforcement=true
     assert.match(output, /enableFloorSettingEnforcement:\s*true/);
   });
 
   it.skip('should update folder floor settings', () => {
-    const output = execSync(
-      `node snippets/updateFolderFloorSettings.js ${folderId}`
-    ).toString();
+    const updateFolderFloorSettings = require('../snippets/updateFolderFloorSettings');
+    const output = updateFolderFloorSettings(folderId).toString();
     // Check that the update was performed
     assert.match(output, /Updated folder floor settings/);
     // Check that the response contains enableFloorSettingEnforcement=true
     assert.match(output, /enableFloorSettingEnforcement:\s*true/);
   });
 
   it.skip('should update project floor settings', () => {
-    const output = execSync(
-      `node snippets/updateProjectFloorSettings.js ${projectId}`
-    ).toString();
+    const updateProjectFloorSettings = require('../snippets/updateProjectFloorSettings');
+    const output = updateProjectFloorSettings(projectId).toString();
     // Check that the update was performed
     assert.match(output, /Updated project floor settings/);
     // Check that the response contains enableFloorSettingEnforcement=true
@@ -999,4 +997,4 @@ describe('Model Armor tests', () => {
       'NO_MATCH_FOUND'
     );
   });
-});
+});