Skip to content

Commit 3cab31c

Browse files
rudrakhsha-crestrudrakhsha-crest
committed
Updated test cases for sanitization snippets, added test for rai filter
1 parent 6fdca28 commit 3cab31c

File tree

1 file changed

+148
-35
lines changed

1 file changed

+148
-35
lines changed

‎model-armor/test/modelarmor.test.js‎

Lines changed: 148 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@ let basicTemplateId;
3434
let basicSdpTemplateId;
3535
let advanceSdpTemplateId;
3636
let templateToDeleteId;
37+
let allFilterTemplateId;
3738
let inspectTemplateName;
3839
let deidentifyTemplateName;
3940

@@ -195,6 +196,8 @@ describe('Model Armor tests', () => {
195196
const SdpAdvancedConfigEnforcement =
196197
protos.google.cloud.modelarmor.v1.SdpAdvancedConfig
197198
.SdpAdvancedConfigEnforcement;
199+
const RaiFilterType =
200+
protos.google.cloud.modelarmor.v1.RaiFilterType;
198201

199202
// Create empty template for sanitizeUserPrompt tests
200203
emptyTemplateId = `${templateIdPrefix}-empty`;
@@ -252,6 +255,41 @@ describe('Model Armor tests', () => {
252255
`projects/${projectId}/locations/${locationId}/templates/${advanceSdpTemplateId}`
253256
);
254257

258+
// Create all-filter template
259+
allFilterTemplateId = `${templateIdPrefix}-all-filters`;
260+
await createTemplate(allFilterTemplateId, {
261+
raiSettings: {
262+
raiFilters: [
263+
{
264+
filterType: RaiFilterType.DANGEROUS,
265+
confidenceLevel: DetectionConfidenceLevel.HIGH,
266+
},
267+
{
268+
filterType: RaiFilterType.HARASSMENT,
269+
confidenceLevel: DetectionConfidenceLevel.HIGH,
270+
},
271+
{
272+
filterType: RaiFilterType.HATE_SPEECH,
273+
confidenceLevel: DetectionConfidenceLevel.HIGH,
274+
},
275+
{
276+
filterType: RaiFilterType.SEXUALLY_EXPLICIT,
277+
confidenceLevel: DetectionConfidenceLevel.HIGH,
278+
},
279+
],
280+
},
281+
piAndJailbreakFilterSettings: {
282+
filterEnforcement: PiAndJailbreakFilterEnforcement.ENABLED,
283+
confidenceLevel: DetectionConfidenceLevel.MEDIUM_AND_ABOVE,
284+
},
285+
maliciousUriFilterSettings: {
286+
filterEnforcement: MaliciousUriFilterEnforcement.ENABLED,
287+
},
288+
});
289+
templatesToDelete.push(
290+
`projects/${projectId}/locations/${locationId}/templates/${allFilterTemplateId}`
291+
);
292+
255293
// Create a template to be deleted
256294
templateToDeleteId = `${templateIdPrefix}-to-delete`;
257295
await createTemplate(templateToDeleteId, {
@@ -273,6 +311,27 @@ describe('Model Armor tests', () => {
273311
await deleteDlpTemplates();
274312
});
275313

314+
// =================== RAI Filter Tests ===================
315+
316+
it('should sanitize user prompt with all RAI filter template', async () => {
317+
const testUserPrompt = "How to make cheesecake without oven at home?";
318+
319+
const output = execSync(
320+
`node snippets/sanitizeUserPrompt.js ${projectId} ${locationId} ${allFilterTemplateId} "${testUserPrompt}"`
321+
).toString();
322+
const response = JSON.parse(output);
323+
324+
assert.equal(
325+
response.sanitizationResult.filterMatchState,
326+
"NO_MATCH_FOUND"
327+
);
328+
329+
assert.equal(
330+
response.sanitizationResult.filterResults.rai.raiFilterResult.matchState,
331+
"NO_MATCH_FOUND"
332+
);
333+
});
334+
276335
// =================== User Prompt Sanitization Tests ===================
277336

278337
it('should detect malicious URL in user prompt', async () => {
@@ -282,11 +341,19 @@ describe('Model Armor tests', () => {
282341
const output = execSync(
283342
`node snippets/sanitizeUserPrompt.js ${projectId} ${locationId} ${basicTemplateId} "${testUserPrompt}"`
284343
).toString();
285-
286-
assert.include(output, '"filterMatchState": "MATCH_FOUND"');
287-
288-
assert.include(output, '"maliciousUriFilterResult"');
289-
assert.include(output, '"matchState": "MATCH_FOUND"');
344+
345+
const response = JSON.parse(output);
346+
347+
assert.equal(
348+
response.sanitizationResult.filterMatchState,
349+
"MATCH_FOUND"
350+
);
351+
352+
assert.equal(
353+
response.sanitizationResult.filterResults.malicious_uris.maliciousUriFilterResult.matchState,
354+
"MATCH_FOUND"
355+
);
356+
290357
assert.include(
291358
output,
292359
'https://testsafebrowsing.appspot.com/s/malware.html'
@@ -299,11 +366,18 @@ describe('Model Armor tests', () => {
299366
const output = execSync(
300367
`node snippets/sanitizeUserPrompt.js ${projectId} ${locationId} ${basicTemplateId} "${testUserPrompt}"`
301368
).toString();
302-
303-
assert.include(output, '"filterMatchState": "NO_MATCH_FOUND"');
304-
305-
assert.include(output, '"csamFilterFilterResult"');
306-
assert.include(output, '"matchState": "NO_MATCH_FOUND"');
369+
370+
const response = JSON.parse(output);
371+
372+
assert.equal(
373+
response.sanitizationResult.filterMatchState,
374+
"NO_MATCH_FOUND"
375+
);
376+
377+
assert.equal(
378+
response.sanitizationResult.filterResults.csam.csamFilterFilterResult.matchState,
379+
"NO_MATCH_FOUND"
380+
);
307381
});
308382

309383
it('should detect jailbreak in user prompt', async () => {
@@ -313,11 +387,18 @@ describe('Model Armor tests', () => {
313387
const output = execSync(
314388
`node snippets/sanitizeUserPrompt.js ${projectId} ${locationId} ${basicTemplateId} "${testUserPrompt}"`
315389
).toString();
316-
317-
assert.include(output, '"filterMatchState": "MATCH_FOUND"');
318-
319-
assert.include(output, '"piAndJailbreakFilterResult"');
320-
assert.include(output, '"matchState": "MATCH_FOUND"');
390+
391+
const response = JSON.parse(output);
392+
393+
assert.equal(
394+
response.sanitizationResult.filterMatchState,
395+
"MATCH_FOUND"
396+
);
397+
398+
assert.equal(
399+
response.sanitizationResult.filterResults.pi_and_jailbreak.piAndJailbreakFilterResult.matchState,
400+
"MATCH_FOUND"
401+
);
321402
});
322403

323404
it('should not detect issues in user prompt with empty template', async () => {
@@ -327,8 +408,13 @@ describe('Model Armor tests', () => {
327408
const output = execSync(
328409
`node snippets/sanitizeUserPrompt.js ${projectId} ${locationId} ${emptyTemplateId} "${testUserPrompt}"`
329410
).toString();
330-
331-
assert.include(output, '"filterMatchState": "NO_MATCH_FOUND"');
411+
412+
const response = JSON.parse(output);
413+
414+
assert.equal(
415+
response.sanitizationResult.filterMatchState,
416+
"NO_MATCH_FOUND"
417+
);
332418
});
333419

334420
// =================== Model Response Sanitization Tests ===================
@@ -340,11 +426,18 @@ describe('Model Armor tests', () => {
340426
const output = execSync(
341427
`node snippets/sanitizeModelResponse.js ${projectId} ${locationId} ${basicTemplateId} "${testModelResponse}"`
342428
).toString();
343-
344-
assert.include(output, '"filterMatchState": "MATCH_FOUND"');
345-
346-
assert.include(output, '"maliciousUriFilterResult"');
347-
assert.include(output, '"matchState": "MATCH_FOUND"');
429+
430+
const response = JSON.parse(output);
431+
432+
assert.equal(
433+
response.sanitizationResult.filterMatchState,
434+
"MATCH_FOUND"
435+
);
436+
437+
assert.equal(
438+
response.sanitizationResult.filterResults.malicious_uris.maliciousUriFilterResult.matchState,
439+
"MATCH_FOUND"
440+
);
348441
});
349442

350443
it('should not detect CSAM in appropriate model response', async () => {
@@ -354,11 +447,18 @@ describe('Model Armor tests', () => {
354447
const output = execSync(
355448
`node snippets/sanitizeModelResponse.js ${projectId} ${locationId} ${basicTemplateId} "${testModelResponse}"`
356449
).toString();
357-
358-
assert.include(output, '"filterMatchState": "NO_MATCH_FOUND"');
359-
360-
assert.include(output, '"csamFilterFilterResult"');
361-
assert.include(output, '"matchState": "NO_MATCH_FOUND"');
450+
451+
const response = JSON.parse(output);
452+
453+
assert.equal(
454+
response.sanitizationResult.filterMatchState,
455+
"NO_MATCH_FOUND"
456+
);
457+
458+
assert.equal(
459+
response.sanitizationResult.filterResults.csam.csamFilterFilterResult.matchState,
460+
"NO_MATCH_FOUND"
461+
);
362462
});
363463

364464
it('should sanitize model response with advanced SDP template', async () => {
@@ -386,8 +486,13 @@ describe('Model Armor tests', () => {
386486
const output = execSync(
387487
`node snippets/sanitizeModelResponse.js ${projectId} ${locationId} ${emptyTemplateId} "${testModelResponse}"`
388488
).toString();
389-
390-
assert.include(output, '"filterMatchState": "NO_MATCH_FOUND"');
489+
490+
const response = JSON.parse(output);
491+
492+
assert.equal(
493+
response.sanitizationResult.filterMatchState,
494+
"NO_MATCH_FOUND"
495+
);
391496
});
392497

393498
it('should detect PII in model response with basic SDP template', async () => {
@@ -402,7 +507,6 @@ describe('Model Armor tests', () => {
402507
assert.include(output, '"sdpFilterResult"');
403508
assert.include(output, '"matchState": "MATCH_FOUND"');
404509

405-
// Check for US_INDIVIDUAL_TAXPAYER_IDENTIFICATION_NUMBER in findings
406510
assert.match(output, /US_INDIVIDUAL_TAXPAYER_IDENTIFICATION_NUMBER/);
407511
});
408512

@@ -417,8 +521,13 @@ describe('Model Armor tests', () => {
417521
const output = execSync(
418522
`node snippets/sanitizeModelResponseWithUserPrompt.js ${projectId} ${locationId} ${emptyTemplateId} "${testUserPrompt}" "${testModelResponse}"`
419523
).toString();
420-
421-
assert.include(output, '"filterMatchState": "NO_MATCH_FOUND"');
524+
525+
const response = JSON.parse(output);
526+
527+
assert.equal(
528+
response.sanitizationResult.filterMatchState,
529+
"NO_MATCH_FOUND"
530+
);
422531
});
423532

424533
it('should sanitize model response with user prompt using advanced SDP template', async () => {
@@ -447,7 +556,11 @@ describe('Model Armor tests', () => {
447556
const output = execSync(
448557
`node snippets/screenPdfFile.js ${projectId} ${locationId} ${basicSdpTemplateId} ${testPdfPath}`
449558
).toString();
450-
451-
assert.include(output, '"filterMatchState": "NO_MATCH_FOUND"');
452-
});
559+
560+
const response = JSON.parse(output);
561+
562+
assert.equal(
563+
response.sanitizationResult.filterMatchState,
564+
"NO_MATCH_FOUND");
565+
});
453566
});

0 commit comments

Comments
 (0)