Skip to content

Commit 8894a90

Browse files
Updated tests as per review comments
1 parent 3cab31c commit 8894a90

File tree

1 file changed

+67
-60
lines changed

1 file changed

+67
-60
lines changed

model-armor/test/modelarmor.test.js

Lines changed: 67 additions & 60 deletions
Original file line numberDiff line numberDiff line change
@@ -196,8 +196,7 @@ describe('Model Armor tests', () => {
196196
const SdpAdvancedConfigEnforcement =
197197
protos.google.cloud.modelarmor.v1.SdpAdvancedConfig
198198
.SdpAdvancedConfigEnforcement;
199-
const RaiFilterType =
200-
protos.google.cloud.modelarmor.v1.RaiFilterType;
199+
const RaiFilterType = protos.google.cloud.modelarmor.v1.RaiFilterType;
201200

202201
// Create empty template for sanitizeUserPrompt tests
203202
emptyTemplateId = `${templateIdPrefix}-empty`;
@@ -314,21 +313,21 @@ describe('Model Armor tests', () => {
314313
// =================== RAI Filter Tests ===================
315314

316315
it('should sanitize user prompt with all RAI filter template', async () => {
317-
const testUserPrompt = "How to make cheesecake without oven at home?";
318-
316+
const testUserPrompt = 'How to make cheesecake without oven at home?';
317+
319318
const output = execSync(
320319
`node snippets/sanitizeUserPrompt.js ${projectId} ${locationId} ${allFilterTemplateId} "${testUserPrompt}"`
321320
).toString();
322321
const response = JSON.parse(output);
323322

324323
assert.equal(
325324
response.sanitizationResult.filterMatchState,
326-
"NO_MATCH_FOUND"
325+
'NO_MATCH_FOUND'
327326
);
328-
327+
329328
assert.equal(
330329
response.sanitizationResult.filterResults.rai.raiFilterResult.matchState,
331-
"NO_MATCH_FOUND"
330+
'NO_MATCH_FOUND'
332331
);
333332
});
334333

@@ -341,19 +340,17 @@ describe('Model Armor tests', () => {
341340
const output = execSync(
342341
`node snippets/sanitizeUserPrompt.js ${projectId} ${locationId} ${basicTemplateId} "${testUserPrompt}"`
343342
).toString();
344-
343+
345344
const response = JSON.parse(output);
346-
347-
assert.equal(
348-
response.sanitizationResult.filterMatchState,
349-
"MATCH_FOUND"
350-
);
351-
345+
346+
assert.equal(response.sanitizationResult.filterMatchState, 'MATCH_FOUND');
347+
352348
assert.equal(
353-
response.sanitizationResult.filterResults.malicious_uris.maliciousUriFilterResult.matchState,
354-
"MATCH_FOUND"
349+
response.sanitizationResult.filterResults.malicious_uris
350+
.maliciousUriFilterResult.matchState,
351+
'MATCH_FOUND'
355352
);
356-
353+
357354
assert.include(
358355
output,
359356
'https://testsafebrowsing.appspot.com/s/malware.html'
@@ -366,17 +363,18 @@ describe('Model Armor tests', () => {
366363
const output = execSync(
367364
`node snippets/sanitizeUserPrompt.js ${projectId} ${locationId} ${basicTemplateId} "${testUserPrompt}"`
368365
).toString();
369-
366+
370367
const response = JSON.parse(output);
371-
368+
372369
assert.equal(
373370
response.sanitizationResult.filterMatchState,
374-
"NO_MATCH_FOUND"
371+
'NO_MATCH_FOUND'
375372
);
376-
373+
377374
assert.equal(
378-
response.sanitizationResult.filterResults.csam.csamFilterFilterResult.matchState,
379-
"NO_MATCH_FOUND"
375+
response.sanitizationResult.filterResults.csam.csamFilterFilterResult
376+
.matchState,
377+
'NO_MATCH_FOUND'
380378
);
381379
});
382380

@@ -387,17 +385,21 @@ describe('Model Armor tests', () => {
387385
const output = execSync(
388386
`node snippets/sanitizeUserPrompt.js ${projectId} ${locationId} ${basicTemplateId} "${testUserPrompt}"`
389387
).toString();
390-
388+
391389
const response = JSON.parse(output);
392-
390+
391+
assert.equal(response.sanitizationResult.filterMatchState, 'MATCH_FOUND');
392+
393393
assert.equal(
394-
response.sanitizationResult.filterMatchState,
395-
"MATCH_FOUND"
394+
response.sanitizationResult.filterResults.pi_and_jailbreak
395+
.piAndJailbreakFilterResult.matchState,
396+
'MATCH_FOUND'
396397
);
397-
398+
398399
assert.equal(
399-
response.sanitizationResult.filterResults.pi_and_jailbreak.piAndJailbreakFilterResult.matchState,
400-
"MATCH_FOUND"
400+
response.sanitizationResult.filterResults.pi_and_jailbreak
401+
.piAndJailbreakFilterResult.confidenceLevel,
402+
'MEDIUM_AND_ABOVE'
401403
);
402404
});
403405

@@ -408,12 +410,12 @@ describe('Model Armor tests', () => {
408410
const output = execSync(
409411
`node snippets/sanitizeUserPrompt.js ${projectId} ${locationId} ${emptyTemplateId} "${testUserPrompt}"`
410412
).toString();
411-
413+
412414
const response = JSON.parse(output);
413-
415+
414416
assert.equal(
415417
response.sanitizationResult.filterMatchState,
416-
"NO_MATCH_FOUND"
418+
'NO_MATCH_FOUND'
417419
);
418420
});
419421

@@ -426,46 +428,50 @@ describe('Model Armor tests', () => {
426428
const output = execSync(
427429
`node snippets/sanitizeModelResponse.js ${projectId} ${locationId} ${basicTemplateId} "${testModelResponse}"`
428430
).toString();
429-
431+
430432
const response = JSON.parse(output);
431-
433+
434+
assert.equal(response.sanitizationResult.filterMatchState, 'MATCH_FOUND');
435+
432436
assert.equal(
433-
response.sanitizationResult.filterMatchState,
434-
"MATCH_FOUND"
437+
response.sanitizationResult.filterResults.malicious_uris
438+
.maliciousUriFilterResult.matchState,
439+
'MATCH_FOUND'
435440
);
436-
441+
437442
assert.equal(
438-
response.sanitizationResult.filterResults.malicious_uris.maliciousUriFilterResult.matchState,
439-
"MATCH_FOUND"
443+
response.sanitizationResult.filterResults.malicious_uris
444+
.maliciousUriFilterResult.maliciousUriMatchedItems[0].uri,
445+
'https://testsafebrowsing.appspot.com/s/malware.html'
440446
);
441447
});
442448

443449
it('should not detect CSAM in appropriate model response', async () => {
444-
const testModelResponse =
445-
'Here is how to teach long division to a child';
450+
const testModelResponse = 'Here is how to teach long division to a child';
446451

447452
const output = execSync(
448453
`node snippets/sanitizeModelResponse.js ${projectId} ${locationId} ${basicTemplateId} "${testModelResponse}"`
449454
).toString();
450-
455+
451456
const response = JSON.parse(output);
452-
457+
453458
assert.equal(
454459
response.sanitizationResult.filterMatchState,
455-
"NO_MATCH_FOUND"
460+
'NO_MATCH_FOUND'
456461
);
457-
462+
458463
assert.equal(
459-
response.sanitizationResult.filterResults.csam.csamFilterFilterResult.matchState,
460-
"NO_MATCH_FOUND"
464+
response.sanitizationResult.filterResults.csam.csamFilterFilterResult
465+
.matchState,
466+
'NO_MATCH_FOUND'
461467
);
462468
});
463469

464470
it('should sanitize model response with advanced SDP template', async () => {
465471
const testModelResponse =
466-
'How can I make my email address test@dot.com make available to public for feedback';
472+
'For following email 1l6Y2@example.com found following associated phone number: 954-321-7890 and this ITIN: 988-86-1234';
467473
const expectedValue =
468-
'How can I make my email address [REDACTED] make available to public for feedback';
474+
'For following email [REDACTED] found following associated phone number: [REDACTED] and this ITIN: [REDACTED]';
469475

470476
const output = execSync(
471477
`node snippets/sanitizeModelResponse.js ${projectId} ${locationId} ${advanceSdpTemplateId} "${testModelResponse}"`
@@ -486,12 +492,12 @@ describe('Model Armor tests', () => {
486492
const output = execSync(
487493
`node snippets/sanitizeModelResponse.js ${projectId} ${locationId} ${emptyTemplateId} "${testModelResponse}"`
488494
).toString();
489-
495+
490496
const response = JSON.parse(output);
491-
497+
492498
assert.equal(
493499
response.sanitizationResult.filterMatchState,
494-
"NO_MATCH_FOUND"
500+
'NO_MATCH_FOUND'
495501
);
496502
});
497503

@@ -521,12 +527,12 @@ describe('Model Armor tests', () => {
521527
const output = execSync(
522528
`node snippets/sanitizeModelResponseWithUserPrompt.js ${projectId} ${locationId} ${emptyTemplateId} "${testUserPrompt}" "${testModelResponse}"`
523529
).toString();
524-
530+
525531
const response = JSON.parse(output);
526-
532+
527533
assert.equal(
528534
response.sanitizationResult.filterMatchState,
529-
"NO_MATCH_FOUND"
535+
'NO_MATCH_FOUND'
530536
);
531537
});
532538

@@ -556,11 +562,12 @@ describe('Model Armor tests', () => {
556562
const output = execSync(
557563
`node snippets/screenPdfFile.js ${projectId} ${locationId} ${basicSdpTemplateId} ${testPdfPath}`
558564
).toString();
559-
565+
560566
const response = JSON.parse(output);
561-
567+
562568
assert.equal(
563569
response.sanitizationResult.filterMatchState,
564-
"NO_MATCH_FOUND");
565-
});
570+
'NO_MATCH_FOUND'
571+
);
572+
});
566573
});

0 commit comments

Comments
 (0)