@@ -168,7 +168,7 @@ class AIGuardInternalTests extends DDSpecification {
168168 return mockResponse(
169169 request,
170170 200 ,
171- [data : [attributes : [action : suite. action, reason : suite. reason, tags : suite. tags ?: [], is_blocking_enabled : suite. blocking]]]
171+ [data : [attributes : [action : suite. action, reason : suite. reason, tags : suite. tags ?: [], tag_probs : suite . tagProbabilities ?: [:], is_blocking_enabled : suite. blocking]]]
172172 )
173173 }
174174 }
@@ -210,12 +210,14 @@ class AIGuardInternalTests extends DDSpecification {
210210 error. action == suite. action
211211 error. reason == suite. reason
212212 error. tags == suite. tags
213+ error. tagProbabilities == suite. tagProbabilities
213214 error. sds == []
214215 } else {
215216 error == null
216217 eval. action == suite. action
217218 eval. reason == suite. reason
218219 eval. tags == suite. tags
220+ eval. tagProbabilities == suite. tagProbabilities
219221 eval. sds == []
220222 }
221223 assertTelemetry(' ai_guard.requests' " action:$suite . action " " block:$throwAbortError " ' error:false'
@@ -555,6 +557,9 @@ class AIGuardInternalTests extends DDSpecification {
555557 if (suite. tags) {
556558 assert meta. attack_categories == suite. tags
557559 }
560+ if (suite. tagProbabilities) {
561+ assert meta. tag_probs == suite. tagProbabilities
562+ }
558563 final receivedMessages = snakeCaseJson(meta. messages)
559564 final expectedMessages = snakeCaseJson(suite. messages)
560565 JSONAssert . assertEquals (expectedMessages, receivedMessages, JSONCompareMode . NON_EXTENSIBLE )
@@ -774,15 +779,17 @@ class AIGuardInternalTests extends DDSpecification {
774779 private final AIGuard.Action action
775780 private final String reason
776781 private final List<String > tags
782+ private final Map<String , Double > tagProbabilities
777783 private final boolean blocking
778784 private final String description
779785 private final String target
780786 private final List<AIGuard.Message > messages
781787
782- TestSuite (AIGuard.Action action , String reason , List <String > tags , boolean blocking , String description , String target , List<AIGuard.Message > messages ) {
788+ TestSuite (AIGuard.Action action , String reason , Map <String , Double > tagProbabilities , boolean blocking , String description , String target , List<AIGuard.Message > messages ) {
783789 this . action = action
784790 this . reason = reason
785- this . tags = tags
791+ this . tags = new ArrayList<> (tagProbabilities. keySet())
792+ this . tagProbabilities = tagProbabilities
786793 this . blocking = blocking
787794 this . description = description
788795 this . target = target
@@ -791,9 +798,9 @@ class AIGuardInternalTests extends DDSpecification {
791798
792799 static List<TestSuite > build () {
793800 def actionValues = [
794- [ALLOW , ' Go ahead'
795- [DENY , ' Nope' ' deny_everything' ' test_deny'
796- [ABORT , ' Kill it with fire' ' alarm_tag' ' abort_everything'
801+ [ALLOW , ' Go ahead' : ]],
802+ [DENY , ' Nope' ' deny_everything' : 0.2D , ' test_deny' : 0.8D ]],
803+ [ABORT , ' Kill it with fire' ' alarm_tag' : 0.1D , ' abort_everything' : 0.9D ]]
797804 ]
798805 def blockingValues = [true , false ]
799806 def suiteValues = [
0 commit comments