@@ -30,7 +30,6 @@ const options = {
3030const client = new ModelArmorClient ( options ) ;
3131const templateIdPrefix = `test-template-${ uuidv4 ( ) . substring ( 0 , 8 ) } ` ;
3232
33- let emptyTemplateId ;
3433let basicTemplateId ;
3534let basicSdpTemplateId ;
3635let advanceSdpTemplateId ;
@@ -261,6 +260,7 @@ describe('Model Armor tests', () => {
261260 before ( async ( ) => {
262261 projectId = await client . getProjectId ( ) ;
263262 const { protos} = require ( '@google-cloud/modelarmor' ) ;
263+
264264 // Import necessary enums
265265 const DetectionConfidenceLevel =
266266 protos . google . cloud . modelarmor . v1 . DetectionConfidenceLevel ;
@@ -275,9 +275,7 @@ describe('Model Armor tests', () => {
275275 . SdpBasicConfigEnforcement ;
276276 const RaiFilterType = protos . google . cloud . modelarmor . v1 . RaiFilterType ;
277277
278- // Create empty template for sanitizeUserPrompt tests
279- emptyTemplateId = `${ templateIdPrefix } -empty` ;
280- await createTemplate ( emptyTemplateId , { } ) ;
278+ await disableFloorSettings ( ) ;
281279
282280 // Create basic template with PI/Jailbreak and Malicious URI filters for sanitizeUserPrompt tests
283281 basicTemplateId = `${ templateIdPrefix } -basic` ;
@@ -360,7 +358,6 @@ describe('Model Armor tests', () => {
360358 } ) ;
361359
362360 templatesToDelete . push (
363- `projects/${ projectId } /locations/${ locationId } /templates/${ emptyTemplateId } ` ,
364361 `projects/${ projectId } /locations/${ locationId } /templates/${ basicTemplateId } ` ,
365362 `projects/${ projectId } /locations/${ locationId } /templates/${ basicSdpTemplateId } ` ,
366363 `projects/${ projectId } /locations/${ locationId } /templates/${ advanceSdpTemplateId } ` ,
@@ -371,10 +368,6 @@ describe('Model Armor tests', () => {
371368
372369 after ( async ( ) => {
373370 for ( const templateName of templatesToDelete ) {
374- // TODO(b/424365799): Uncomment below code once the mentioned issue is resolved
375- // Disable floor settings to restore original state
376- // await disableFloorSettings();
377-
378371 await deleteTemplate ( templateName ) ;
379372 }
380373
@@ -728,24 +721,6 @@ describe('Model Armor tests', () => {
728721 ) ;
729722 } ) ;
730723
731- it ( 'should not detect issues in user prompt with empty template' , async ( ) => {
732- const sanitizeUserPrompt = require ( '../snippets/sanitizeUserPrompt' ) ;
733- const testUserPrompt =
734- 'Can you describe this link? https://testsafebrowsing.appspot.com/s/malware.html,' ;
735-
736- const response = await sanitizeUserPrompt (
737- projectId ,
738- locationId ,
739- emptyTemplateId ,
740- testUserPrompt
741- ) ;
742-
743- assert . equal (
744- response . sanitizationResult . filterMatchState ,
745- 'NO_MATCH_FOUND'
746- ) ;
747- } ) ;
748-
749724 // =================== Model Response Sanitization Tests ===================
750725
751726 it ( 'should detect malicious URL in model response' , async ( ) => {
@@ -826,24 +801,6 @@ describe('Model Armor tests', () => {
826801 ) ;
827802 } ) ;
828803
829- it ( 'should not detect issues in model response with empty template' , async ( ) => {
830- const sanitizeModelResponse = require ( '../snippets/sanitizeModelResponse' ) ;
831- const testModelResponse =
832- 'For following email 1l6Y2@example.com found following associated phone number: 954-321-7890 and this ITIN: 988-86-1234' ;
833-
834- const response = await sanitizeModelResponse (
835- projectId ,
836- locationId ,
837- emptyTemplateId ,
838- testModelResponse
839- ) ;
840-
841- assert . equal (
842- response . sanitizationResult . filterMatchState ,
843- 'NO_MATCH_FOUND'
844- ) ;
845- } ) ;
846-
847804 it ( 'should detect PII in model response with basic SDP template' , async ( ) => {
848805 const sanitizeModelResponse = require ( '../snippets/sanitizeModelResponse' ) ;
849806 const testModelResponse =
@@ -874,27 +831,6 @@ describe('Model Armor tests', () => {
874831
875832 // =================== Model Response with User Prompt Tests ===================
876833
877- it ( 'should not detect issues in model response with user prompt using empty template' , async ( ) => {
878- const sanitizeModelResponseWithUserPrompt = require ( '../snippets/sanitizeModelResponseWithUserPrompt' ) ;
879- const testUserPrompt =
880- 'How can I make my email address test@dot.com make available to public for feedback' ;
881- const testModelResponse =
882- 'You can make support email such as contact@email.com for getting feedback from your customer' ;
883-
884- const response = await sanitizeModelResponseWithUserPrompt (
885- projectId ,
886- locationId ,
887- emptyTemplateId ,
888- testModelResponse ,
889- testUserPrompt
890- ) ;
891-
892- assert . equal (
893- response . sanitizationResult . filterMatchState ,
894- 'NO_MATCH_FOUND'
895- ) ;
896- } ) ;
897-
898834 it ( 'should sanitize model response with user prompt using advanced SDP template' , async ( ) => {
899835 const sanitizeModelResponseWithUserPrompt = require ( '../snippets/sanitizeModelResponseWithUserPrompt' ) ;
900836 const testUserPrompt =
@@ -937,57 +873,66 @@ describe('Model Armor tests', () => {
937873 'NO_MATCH_FOUND'
938874 ) ;
939875 } ) ;
876+ } ) ;
877+
878+ describe ( 'Model Armor floor setting tests' , ( ) => {
879+ before ( async ( ) => {
880+ projectId = await client . getProjectId ( ) ;
881+ } ) ;
882+
883+ after ( async ( ) => {
884+ await disableFloorSettings ( ) ;
885+ } ) ;
940886
941- // =================== Floor Settings Tests ===================
887+ it ( 'should update organization floor settings' , async ( ) => {
888+ const updateOrganizationFloorSettings = require ( '../snippets/updateOrganizationFloorSettings' ) ;
889+ const output = await updateOrganizationFloorSettings . main ( organizationId ) ;
890+ // Check that the enableFloorSettingEnforcement=true
891+ assert . equal ( output . enableFloorSettingEnforcement , true ) ;
892+ } ) ;
942893
943- // TODO(b/424365799): Enable below tests once the mentioned issue is resolved
894+ it ( 'should update folder floor settings' , async ( ) => {
895+ const updateFolderFloorSettings = require ( '../snippets/updateFolderFloorSettings' ) ;
896+ const output = await updateFolderFloorSettings . main ( folderId ) ;
897+ // Check that the enableFloorSettingEnforcement=true
898+ assert . equal ( output . enableFloorSettingEnforcement , true ) ;
899+ } ) ;
944900
945- it . skip ( 'should get organization floor settings' , async ( ) => {
901+ it ( 'should update project floor settings' , async ( ) => {
902+ const updateProjectFloorSettings = require ( '../snippets/updateProjectFloorSettings' ) ;
903+ const output = await updateProjectFloorSettings . main ( projectId ) ;
904+ // Check that the enableFloorSettingEnforcement=true
905+ assert . equal ( output . enableFloorSettingEnforcement , true ) ;
906+ } ) ;
907+
908+ it ( 'should get organization floor settings' , async ( ) => {
946909 const getOrganizationFloorSettings = require ( '../snippets/getOrganizationFloorSettings' ) ;
947910
948911 const output = await getOrganizationFloorSettings . main ( organizationId ) ;
949912
950913 const expectedName = `organizations/${ organizationId } /locations/global/floorSetting` ;
951914 assert . equal ( output . name , expectedName ) ;
915+ assert . exists ( output . enableFloorSettingEnforcement ) ;
952916 } ) ;
953917
954- it . skip ( 'should get folder floor settings' , async ( ) => {
918+ it ( 'should get folder floor settings' , async ( ) => {
955919 const getFolderFloorSettings = require ( '../snippets/getFolderFloorSettings' ) ;
956920
957921 const output = await getFolderFloorSettings . main ( folderId ) ;
958922
959923 // Check for expected name format in output
960924 const expectedName = `folders/${ folderId } /locations/global/floorSetting` ;
961925 assert . equal ( output . name , expectedName ) ;
926+ assert . exists ( output . enableFloorSettingEnforcement ) ;
962927 } ) ;
963928
964- it . skip ( 'should get project floor settings' , async ( ) => {
929+ it ( 'should get project floor settings' , async ( ) => {
965930 const getProjectFloorSettings = require ( '../snippets/getProjectFloorSettings' ) ;
966931
967932 const output = await getProjectFloorSettings . main ( projectId ) ;
968933 // Check for expected name format in output
969934 const expectedName = `projects/${ projectId } /locations/global/floorSetting` ;
970935 assert . equal ( output . name , expectedName ) ;
971- } ) ;
972-
973- it . skip ( 'should update organization floor settings' , async ( ) => {
974- const updateOrganizationFloorSettings = require ( '../snippets/updateOrganizationFloorSettings' ) ;
975- const output = await updateOrganizationFloorSettings . main ( organizationId ) ;
976- // Check that the enableFloorSettingEnforcement=true
977- assert . equal ( output . enableFloorSettingEnforcement , true ) ;
978- } ) ;
979-
980- it . skip ( 'should update folder floor settings' , async ( ) => {
981- const updateFolderFloorSettings = require ( '../snippets/updateFolderFloorSettings' ) ;
982- const output = await updateFolderFloorSettings . main ( folderId ) ;
983- // Check that the enableFloorSettingEnforcement=true
984- assert . equal ( output . enableFloorSettingEnforcement , true ) ;
985- } ) ;
986-
987- it . skip ( 'should update project floor settings' , async ( ) => {
988- const updateProjectFloorSettings = require ( '../snippets/updateProjectFloorSettings' ) ;
989- const output = await updateProjectFloorSettings . main ( projectId ) ;
990- // Check that the enableFloorSettingEnforcement=true
991- assert . equal ( output . enableFloorSettingEnforcement , true ) ;
936+ assert . exists ( output . enableFloorSettingEnforcement ) ;
992937 } ) ;
993938} ) ;
0 commit comments